Sunbelt Computer Software

PL/B Language Development and Support

python/data_structures/binary_tree/binary_tree_node_sum.py at master · bchetty/python · GitHub

75 lines (56 loc) · 1.71 KB

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

"""

Sum of all nodes in a binary tree.

Python implementation:

O(n) time complexity - Recurses through :meth:`depth_first_search`

with each element.

O(n) space complexity - At any point in time maximum number of stack

frames that could be in memory is `n`

"""

from __future__ import annotations

from collections.abc import Iterator

class Node:

"""

A Node has a value variable and pointers to Nodes to its left and right.

"""

def __init__(self, value: int) -> None:

self.value = value

self.left: Node | None = None

self.right: Node | None = None

class BinaryTreeNodeSum:

r"""

The below tree looks like this

10

/ \

5 -3

/ / \

12 8 0

>>> tree = Node(10)

>>> sum(BinaryTreeNodeSum(tree))

10

>>> tree.left = Node(5)

>>> sum(BinaryTreeNodeSum(tree))

15

>>> tree.right = Node(-3)

>>> sum(BinaryTreeNodeSum(tree))

12

>>> tree.left.left = Node(12)

>>> sum(BinaryTreeNodeSum(tree))

24

>>> tree.right.left = Node(8)

>>> tree.right.right = Node(0)

>>> sum(BinaryTreeNodeSum(tree))

32

"""

def __init__(self, tree: Node) -> None:

self.tree = tree

def depth_first_search(self, node: Node | None) -> int:

if node is None:

return 0

return node.value + (

self.depth_first_search(node.left) + self.depth_first_search(node.right)

)

def __iter__(self) -> Iterator[int]:

yield self.depth_first_search(self.tree)

if __name__ == "__main__":

import doctest

doctest.testmod()