[PYTHON] [At Coder] Acing C-XYZ Triplets

This is the C problem of the Acing Programming Contest 2020. I had a hard time, so I looked back with the code that failed.

Failure 1 (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))

This is a method of defining a function f_count (j) that finds f (n), assigning 1 to n to j, and outputting.

Since the time was over, I tried erasing z, but it was still over time.

** <Point 1> **

Looking at the answer, the maximum value of n is 10 ** 4. Since x, y, z are squared, the maximum value of x, y, z is less than 100.

** <Point 2> **

At f (1), search x, y, z from 1 to 100. → At f (2), search x, y, z from 1 to 100. It takes time.

Failure 2

N=int(input())
#f(N)List to put
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

Substitute from 1 to 100 for x, y, z, calculate the value of x ** 2 + y ** 2 + z ** 2 + x * y + y * z + z * x, and list the result ans_list I put it in.

However, I made a mistake in the range. Since ans_list [0] = [], ans_list [1] = f (1), ans_list [N-1] = f (N-1), f (N) is no longer included in the list.

Correct answer 1

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])

I set ans_list [s-1] + = 1.

Correct answer 2

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)

For correct answer 1, I tried to make it a dictionary because the index of the list and n of f (n) were not clear.