Togetherness of a family PoD solution

code.py

import math
# Number of people
n = int(input())
tokens = [] # For storing tokens
tree = [] # For storing adjacency list in tree, check https://youtu.be/09zltCNaRF8 for more
result = [] # For storing the results


#Slightly improved logic to check if a number is prime or not
def isPrime(number):
    if(number <= 1): return False
    if(number <= 3): return True

    if(number % 2 == 0 or number % 3 == 0):
        return False
    
    for i in range(5, math.ceil(math.sqrt(number)) + 1, 6):
        if(number % i == 0 or (number % (i + 2) == 0)):
            return False

    return True

# Depth first traversal as explained:https://youtu.be/pEwTw97BS2c?t=340
def depth_first_traversal(source):
    for child in tree[source]:
        #Go deep into the branch
        depth_first_traversal(child)
    #While coming up, bubble up the results from children to the parent
    for child in tree[source]:
        if(isPrime(tokens[child-1])):
            result[source] += 1
        result[source] += result[child]

#Get all the tokens
tokens = list(map(int, input().split()))
#Initialise tree and results
for i in range(n+1):
    tree.append([])
    result.append(0)
#Store the tree as adjacency list
for i in range(n-1):
    x,y = map(int, input().split())
    tree[x].append(y)
#Perform depth first traversal and build the results list,
#this operation only needs to run once
depth_first_traversal(1)
#Answer all the queries
q = int(input())
for _ in range(q):
    x = int(input())
    print(result[x])