CSES - Syvyydet

Tehtäväsi on laskea, montako solmua annetussa puussa on tietyllä syvyydellä.

Esimerkiksi alla olevassa puussa solmu 11 on syvyydellä 00, solmut 44, 55 ja 22 ovat syvyydellä 11 sekä solmut 33, 77 ja 66 ovat syvyydellä 22.

Toteuta tiedostoon depths.py funktio count_nodes, jolle annetaan parametrina viittaus puun juurisolmuun sekä haluttu solmun syvyys. Funktio palauttaa solmujen määrän.

class Node:
    def __init__(self, value, children=None):
        self.value = value
        self.children = children if children else []

    def __repr__(self):
        if self.children == []:
            return f"Node({self.value})"
        else:
            return f"Node({self.value}, {self.children})"

def count_nodes(node, depth):
    # TODO

if __name__ == "__main__":
    tree1 = Node(1, [Node(4, [Node(3), Node(7)]),
                     Node(5),
                     Node(2, [Node(6)])])
    print(count_nodes(tree1, 0)) # 1
    print(count_nodes(tree1, 1)) # 3
    print(count_nodes(tree1, 2)) # 3
    print(count_nodes(tree1, 3)) # 0

    tree2 = Node(1, [Node(2, [Node(3, [Node(4)])])])
    print(count_nodes(tree2, 0)) # 1
    print(count_nodes(tree2, 1)) # 1
    print(count_nodes(tree2, 2)) # 1
    print(count_nodes(tree2, 3)) # 1
    print(count_nodes(tree2, 4)) # 0

    tree3 = Node(1, [Node(2), Node(3), Node(4)])
    print(count_nodes(tree3, 0)) # 1
    print(count_nodes(tree3, 1)) # 3
    print(count_nodes(tree3, 2)) # 0