0674 - Count 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.

Date de intrare

Fișierul de intrare countsubin.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ătoarele k 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 countsubout.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

countsubin.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

countsubout.txt

3

2

6

2

Exemplul 2

countsubin.txt

5

4 2 3

7 4 0

1 0 0

5 0 0

2 0 0

3

1

2

4

countsubout.txt

2

2

2

Rezolvare

<syntaxhighlight lang="python" line>

1. 0674 - Count Sub

class TreeNode:

   def __init__(self, value, left=None, right=None):
       self.value = value
       self.left = left
       self.right = right

def read_tree_from_file(file):

   n = int(file.readline())
   if not (1 <= n <= 1000):
       return None

   nodes = [None]  # index 0 nu este folosit

   for i in range(1, n + 1):
       x, st, dr = map(int, file.readline().split())
       if not (1 <= x <= 1000000):
           return None
       if not (0 <= st <= n):
           return None
       if not (0 <= dr <= n):
           return None
       nodes.append(TreeNode(x, st, dr))

   return nodes

def count_subtree_nodes(node, nodes):

   if node == 0:
       return 0
   return 1 + count_subtree_nodes(nodes[node].left, nodes) + count_subtree_nodes(nodes[node].right, nodes)

def check_restrictions(k):

   return 1 <= k <= 1000

def main():

   with open('countsubin.txt', 'r') as file:
       nodes = read_tree_from_file(file)
       if nodes is None:
           print("false")
           return

       n = len(nodes) - 1  # Numărul de noduri

       k = int(file.readline())
       if not check_restrictions(k):
           print("false")
           return

       results = []

       for _ in range(k):
           query_node = int(file.readline())
           if not (1 <= query_node <= n):
               print("false")
               return
           subtree_nodes = count_subtree_nodes(query_node, nodes)
           results.append(subtree_nodes)

   with open('countsubout.txt', 'w') as file:
       for result in results:
           file.write(str(result) + '\n')

if __name__ == "__main__":

   main()

</syntaxhighlight>

 Explicatie 

Exemplul corespunde arborelui de mai jos, în care au fost marcate cu albastru valorile din noduri, iar cu roșu numerele de ordine ale nodurilor.