Atcoder エイシングプログラミングコンテスト Python (A~C)

総括

難しかったです。A,B,Cの3問ACでした。
2進数がよくわからないのでDは問題をみた瞬間スキップ、E問題はできそうな気がしたものの回答例をみてもよくわからず・・・。

問題

(https://atcoder.jp/contests/aising2020/tasks)

A. Number of Multiples

回答

L, R, d = map(int, input().split())

count = 0
for i in range(L, R+1):
    if i % d == 0:
        count += 1

print(count)

素直に書きました。ほかの方はもっとかっこいい書き方をしているみたいですが。

B. An Odd Problem

回答

N = int(input())
a = list(map(int, input().split()))

count = 0
for i in range(0, N, 2):
    if i % 2 == 0 and a[i] % 2 != 0:
        count += 1
print(count)

こちらも素直に書きました。
問題分は奇数ですがリストの添え字は偶数をとってきます。

C. XYZ Triplets

回答

N = int(input())

keys = [i for i in range(1, N+1)]
values = [0] * N
count_dict = dict(zip(keys, values))

for x in range(1, N//6+7):
    for y in range(x, N//6+7):
        for z in range(y, N//6+7):
            n = x**2 + y**2 + z**2 + x*y + y*z + z*x
            if n > N:
                break
            if x == y and y == z:
                count_dict[n] += 1
            elif x != y and x != z and y != z:
                count_dict[n] += 6
            else:
                count_dict[n] += 3

for v in count_dict.values():
    print(v)

素直に書くとx,y,zそれぞれでrange(1, N+1)のforループを回してx**2 + y**2 + z**2 + x*y + y*z + z*x = nを判定すればよさそうです。
しかし制約が10**4なので普通に3重ループをすると間に合いませんので計算量を抑えることを考えます。

まず考えたのは、(x,y,z) = (1,1,1)である場合nは6なので、どれだけ多くてもforループで回す最大値はN//6くらいでよさそうです。ただし添え字を+1しなければいけないことを考えると余裕をもってN//6+7としました。
(本番ではN//6+7を使用しましたが、10**4という制約を逆に使うと、(x,y,z)は高々100くらいまでしか使えないということがわかりますね。これは本番中に気づきませんでした・・・。)

で、これだけではたぶんぎりぎり間に合わなそうなので、もうちょっと計算量を減らすことを考えます。

何通りか例を挙げてみると下記のことに気が付きます。
- (1,2,3)と(1,3,2)と(2,1,3)と(2,3,1)と(3,1,2)と(3,2,1)のようにx,y,zがすべて違う場合にnが同値になる組み合わせは6つある
- (1,1,2)と(1,2,1)と(2,1,1)のように2つが同じで1つが違う場合にnが同値になる組み合わせは3つある
- (1,1,1)のようにすべて同じ数字の場合はnが同値になる組み合わせは1つある

したがって、x,y,zのforループの範囲をそれぞれrange(1, N//6+7)、range(x, N//6+7)、range(y, N//6+7)として、上記の3通りの場合でcountを+6、+3、+1してやればよさそうです。

あとはnがNを超えた場合にbreakいれておけば何とか間に合います。