I have a problem with my code, I use the eratosthenes this function to create a list, and I just make a big memory size list (1000000), in the main I input the codeforces input() and let them sqrt(), so, I dont know why my code timelimit over than 2000ms, please help me to solve this problem.
import math
def eratosthenes(n): # creative the prime list
IsPrime = [True] * (n + 1)
IsPrime[1] = False
for i in range(2, int(n ** 0.5) + 1):
if IsPrime[i]:
for j in range(i * i, n + 1, i):
IsPrime[j] = False
return [x for x in range(2, n + 1) if IsPrime[x]]
# main code
input()
li = list(map(int, input().split()))
nl = eratosthenes(1000000)
for i in li:
i = math.sqrt(i)
if int(i) == i:
print("YES" if i in nl else "NO")
else:
print("NO")
I think the run time of my code is similar to other's code when I refer them, or I'm wrong.