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

#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()

 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.