I have don this problem in py3 and all the logic part is correct and while running local run time = 600ms codeforce runtime TLE

can any one help me with this my code --> """ import math

def callone(num): def prime(n): prime_flag = 0 if(n > 2): for i in range(2, int(math.sqrt(n)) + 1): if (n % i == 0): prime_flag = 1 break if (prime_flag == 0): return 1 else: return 0 else: return 0

if(num > 4):
    x = math.sqrt(num)
    s = str(x)
    if(s[-2:] == '.0'):
        if(prime(x)):
                return 1
        else:
            return 0
    else:
        return 0
else:
    return 0

def array(num): ele = 3 primearr = [] while(ele < num+1): for i in range(2, int(math.sqrt(ele)) + 1): if (ele % i == 0): break else: primearr.append(ele) ele+=1 return(primearr)

def call(num): if(num > 4): x = math.sqrt(num) s = str(x) if(s[-2:] == '.0'): return 1 else: return 0 else: return 0

n = int(input()) arr = list(map(int, input().split()))

if len(arr) > 1: m = int(math.sqrt(max(arr))) primearr = array(m) # print(primearr) for i in arr: if(i==4): print('YES') elif(i%2==0): print('NO') elif(call(i) and (int(math.sqrt(i)) in primearr)): print("YES") else: print("NO") else: if(callone(arr[0])): print("YES") else: print("NO")

"""