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

<syntaxhighlight lang="python" line> 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()

</syntaxhighlight>