【AtCoder】エイシング C - XYZ Triplets

エイシング プログラミング コンテスト 2020のC問題です。
苦戦したので、失敗したコードと共に振り返ってみました。

失敗１(TLE)

n=int(input())
def f_count(j):
   count=0
   for x in range(1,j):
       for y in range(1,j):
           for z in range(1,j):
               if x**2+y**2+z**2+x*y+y*z+z*x==j:
                   count+=1
   return count

for i in range(1,n+1):
   print(f_count(i))

f(n)を求める関数f_count(j)を定義して、jに1～nまで代入して、出力する方法です。

時間超過になったので、zを消したりしてみましたが、それでも時間超過になりました。

＜ポイント１＞

解答を見てみると、nの最大値は10**4。x,y,zは二乗しているから、x,y,zの最大値は100より小さい。

＜ポイント２＞

f(1)の時、x,y,zを1～100まで全探索する → f(2)の時、x,y,zを1～100まで全探索する…としていては、時間がかかる。

失敗２

N=int(input())
#f(N)を入れるリスト
ans_list=[0]*N
for x in range(1,100):
   for y in range(1,100):
       for z in range(1,100):
           s=x**2+y**2+z**2+x*y+y*z+z*x
           if s<=N:
               ans_list[s]+=1


for i in range(1,N+1):
  print(ans_list[i])

＜結果＞IndexError: list index out of range

x, y, zに1～100まで代入して、x*2+y2+z*2+x*y+y*z+z*xの値を計算し、結果をリストans_listに入れていきました。

しかし、範囲を間違えました。ans_list[0]=[], ans_list[1]=f(1), ans_list[N-1]=f(N-1)としてしまったため、f(N)がリストに入らなくなってしまいました。

正解１

N=int(input())
max_x=int(N**0.5)
ans_list=[0]*N

for x in range(1,max_x):
   for y in range(1,max_x):
       for z in range(1,max_x):
           s=x**2+y**2+z**2+x*y+y*z+z*x
           if s<=N:
               ans_list[s-1]+=1

for i in range(N):
   print(ans_list[i])

ans_list[s-1]+=1にしました。​

正解２

n=int(input())

keys=[i for i in range(1,n+1)]
values=[0 for i in range(1,n+1)]
dic=dict(zip(keys,values))

for x in range(1,100):
    for y in range(1,100):
        for z in range(1,100):
            num=x**2+y**2+z**2+x*y+y*z+z*x
            if 1<=num<=n:
                dic[num]+=1


for key in dic:
    value=dic[key]
    print(value)

正解１は、リストのインデックスとf(n)のnがズレるのがすっきりしなかったので、辞書にしてみました。