0676 - Count Prim Sub

Cerinţa

Considerăm un arbore binar cu n noduri în care fiecare nod este numerotat de la 1 la n și conține o valoare număr natural. Se dau k noduri din arbore și se cere determinarea, pentru fiecare nod, a numărului de noduri din subarborele cu rădăcina în acel nod care conțin valori prime.

Date de intrare

Fișierul de intrare countprimsubin.txt conține pe prima linie numărul n. Fiecare dintre următoarele n linii contine câte 3 numere X st dr; linia i + 1 din fișier conține informatiile despre nodul numerotat cu i: X reprezintă valoare din nod, st reprezintă numărul de ordine al descendentului stâng sau 0 dacă nodul i nu are descendent stâng, iar dr reprezintă numărul de ordine al descendentului drept sau 0 dacă nodul i nu are descendent drept.

Pe următoarea linie se află numărul k, iar pe fiecare dintre următoarelek linii se află câte un număr natural cuprins între 1 și n, reprezentând nodul curent.

Date de ieșire

Fișierul de ieșire countprimsubout.txt va conține k linii; fiecare linie va conține numărul de noduri din subarborele cu rădăcina în nodul corespunzător.

Restricţii şi precizări

• 1 ≤ n ≤ 1000
• valorile din nodurile arborelui vor fi mai mici sau egale cu 1.000.000
• 1 ≤ k ≤ 1000

Exemplul 1

countprimsubin.txt

6
2 3 5
6 0 6
1 0 0
7 1 2
4 0 0
10 0 0
4
1
2
4
2

countprimsubout.txt

Datele de intrare corespund restrictiilor impuse
1
0
2
0

Exemplul 2

countprimsubin.txt

1002
2 0 0
3 0 0
5 0 0
7 0 0
1000001 0 0
2
1
5

countprimsubout.txt

Datele de intrare nu corespund restrictiilor impuse

Rezolvare

tree = []
prime_count = []


def is_prime(n):
    if n < 2:
        return False
    if n == 2 or n == 3:
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False
    i = 5
    w = 2
    while i * i <= n:
        if n % i == 0:
            return False
        i += w
        w = 6 - w
    return True


def dfs(node):
    global tree, prime_count
    prime_count[node] = is_prime(tree[node][0])
    if tree[node][1] != 0:
        prime_count[node] += dfs(tree[node][1])
    if tree[node][2] != 0:
        prime_count[node] += dfs(tree[node][2])
    return prime_count[node]


def main():
    global tree, prime_count
    with open('countprimsubin.txt', 'r') as fin, open('countprimsubout.txt', 'w') as fout:
        n = int(fin.readline().strip())
        tree = [[0, 0, 0] for _ in range(n+1)]
        prime_count = [0 for _ in range(n+1)]
        for i in range(1, n+1):
            x, st, dr = map(int, fin.readline().strip().split())
            tree[i] = [x, st, dr]
        k = int(fin.readline().strip())
        nodes = [int(fin.readline().strip()) for _ in range(k)]
        if n < 1 or n > 1000 or any(x < 1 or x > 1000000 for x, _, _ in tree[1:]) or k < 1 or k > 1000:
            fout.write("Datele de intrare nu corespund restrictiilor impuse\n")
            return
        fout.write("Datele de intrare corespund restrictiilor impuse\n")
        for node in nodes:
            fout.write(str(dfs(node)) + '\n')


if __name__ == "__main__":
    main()