N-Tree Sum

tree_node_sum.py

#!/usr/bin/env python
from __future__ import annotations
from uuid import uuid4, UUID
from typing import Sequence


class Node:
    node_id: UUID
    time_to_complete: int
    children: Sequence[Node]

    def __init__(self, time_to_complete: int, children: Sequence[Node] = []) -> None:
        self.node_id = uuid4()
        self.children = set(children)
        self.time_to_complete = time_to_complete

    def __hash__(self) -> int:
        return hash(f"{self.node_id}:{self.time_to_complete}")

    def __gt__(self, x: Node) -> bool:
        try:
            return self.time_to_complete > x.time_to_complete
        except ValueError:
            return self.time_to_complete > x

    def __ge__(self, x: Node) -> bool:
        try:
            return self.time_to_complete >= x.time_to_complete
        except ValueError:
            return self.time_to_complete >= x

    def __lt__(self, x: Node) -> bool:
        try:
            return self.time_to_complete < x.time_to_complete
        except ValueError:
            return self.time_to_complete < x

    def __le__(self, x: Node) -> bool:
        try:
            return self.time_to_complete <= x.time_to_complete
        except ValueError:
            return self.time_to_complete <= x

    def sum(self) -> int:
        ttc = self.time_to_complete
        for child in self.children:
            ttc += child.sum()
        return ttc


if __name__ == "__main__":
    D = Node(time_to_complete=4)
    C = Node(time_to_complete=3)
    B = Node(time_to_complete=2, children=[C, D])
    A = Node(time_to_complete=1, children=[B])
    print(A.sum())