0760 - k Nivel 1

Cerinţa

Se consideră un arbore binar în care nodurile memorează numere naturale nenule și un număr k. În arbore rădăcina se află pe nivelul 0, fii rădăcinii pe nivelul 1, fii fiilor rădăcinii pe nivelul 2, etc. Să se determine suma valorilor din nodurile aflate pe nivelul k.

Date de intrare

Fișierul de intrare knivel1in.txt conține pe prima linie lista valorilor memorate în nodurile arborelui, obținute în urma parcurgerii în preordine (rădăcină, stâng, drept). Dacă un nod nu are descendent stâng, în listă va apare valoarea 0. Dacă un nod nu are descendent drept, în listă va apare valoarea 0.

Linia a doua a fișierului conține numărul k.

Date de ieșire

Fișierul de ieșire knivel1out.txt va conține pe prima linie un singur număr S, reprezentând suma cerută

Restricţii şi precizări

• se recomandă folosirea arborilor alocați dinamic.

Exemplul 1

knivel1in.txt

67 51 18 0 0 24 0 0 48 0 11 0 0
1

knivel1out.txt

Datele de intrare corespund restrictiilor impuse
99

Exemplul 2

knivel1in.txt

1 0 0 
1 

knivel1out.txt

Datele de intrare nu corespund restrictiilor impuse

Rezolvare

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


def build_tree(values):
    value = next(values)
    if value == 0:
        return None
    node = Node(value)
    node.left = build_tree(values)
    node.right = build_tree(values)
    return node


def sum_at_level(node, k):
    if node is None:
        return 0
    if k == 0:
        return node.value
    return sum_at_level(node.left, k-1) + sum_at_level(node.right, k-1)


def main():
    with open('knivel1in.txt', 'r') as fin:
        values = list(map(int, fin.readline().split()))
        k = int(fin.readline())
    root = build_tree(iter(values))
    if root is None or root.left is None or root.right is None:
        print("Datele de intrare nu corespund restrictiilor impuse")
        return
    print("Datele de intrare corespund restrictiilor impuse")
    s = sum_at_level(root, k)
    with open('knivel1out.txt', 'w') as fout:
        fout.write(str(s))


if __name__ == "__main__":
    main()