tree
β’ 615 words β’ 3 min β’ updated
Problem statement: Implement a glorified Node class, representing a binary
tree.
Employ:
- type annotations
- BFS, DFS
- pre-order, in-order, post-order traversals
- height
- length / size / count
- insertion
- clear
- root, leaves
- map, reduce
python
from collections import deque
from typing import TypeVar, Generic, Optional, Any
from operator import __add__
T = TypeVar('T')
class Node(Generic[T]):
## def __init__(self, val: T | None = None):
def __init__(self, val: Optional[T] = None):
self.val = val
self.left = None
self.right = None
## def get(self) -> T | None:
def get(self) -> Optional[T]:
return self.val
def __len__(self) -> int:
return sum((
self.val is not None,
len(self.left) if self.left else 0,
len(self.right) if self.right else 0,
))
def insert_left(self, node: 'Node[T]') -> None:
self.left = node
def insert_right(self, node: 'Node[T]') -> None:
self.right = node
def height(self) -> int:
return 1 + max(
self.left.height() if self.left else 0,
self.right.height() if self.right else 0,
)
def dfs(self) -> list[T]:
ans = []
def dfs(node) -> None:
if not node:
return
ans.append(node.val)
dfs(node.left)
dfs(node.right)
dfs(self)
return ans
def bfs(self) -> list[T]:
ans = []
queue = deque([self])
while queue:
node = queue.popleft()
ans.append(node.val)
for child in [node.left, node.right]:
if child:
queue.append(child)
return ans
def sum(self) -> int:
## return self.reduce(lambda x, y: x + y, 0)
## return self.reduce(__add__, 0)
def dfs(node) -> int:
if not node:
return 0
return node.val + dfs(node.left) + dfs(node.right)
return dfs(self)
def __eq__(self, other: object) -> bool:
if not isinstance(other, Node):
return False
if self.val != other.val:
return False
return self.left == other.left and self.right == other.right
def preorder(self) -> list[int]:
ans = []
def preorder(node):
if node:
ans.append(node.val)
preorder(node.left)
preorder(node.right)
preorder(self)
return ans
def inorder(self) -> list[int]:
ans = []
def inorder(node):
if node:
inorder(node.left)
ans.append(node.val)
inorder(node.right)
inorder(self)
return ans
def postorder(self) -> list[int]:
ans = []
def postorder(node):
if node:
postorder(node.left)
postorder(node.right)
ans.append(node.val)
postorder(self)
return ans
def root(self) -> 'Node[T]':
return self
def leaves(self) -> list[T]:
## return [node.val for node in self.dfs() if node is not None]
ans = []
def dfs(node):
if not node:
return
if not node.left and not node.right:
ans.append(node.val)
dfs(node.left)
dfs(node.right)
dfs(self)
return ans
def map(self, f) -> None:
def dfs(node):
if node:
node.val = f(node.val)
dfs(node.left)
dfs(node.right)
dfs(self)
def reduce(self, f, init: int = 0) -> Any:
ans = init
stack = [self]
while stack:
node = stack.pop()
ans = f(ans, node.val)
if node.left:
stack.append(node.left)
if node.right:
stack.append(node.right)
return ans
def clear(self) -> None:
## shallow clear
## self.val = None
## self.left = None
## self.right = None
# deep clear
def dfs(node):
if node:
node.val = None
dfs(node.left)
dfs(node.right)
node.left = None
node.right = None
dfs(self)
if __name__ == "__main__":
t = Node()
# XXX: assert t
assert t.get() is None
assert Node(5).get() == 5
assert len(Node(5)) == 1
t = Node(1)
n2 = Node(2)
t.insert_left(n2)
t.insert_right(Node(3))
n2.insert_left(Node(4))
"""
1 -> 2 -> 4
|-> 3
"""
assert len(t) == 4
assert t.height() == 3
assert t.sum() == sum(range(1, 5))
assert t.dfs() == [1, 2, 4, 3]
assert t.bfs() == [1, 2, 3, 4]
assert set(t.dfs()) == set(t.bfs())
w = Node(1)
n2 = Node(2)
w.insert_left(n2)
w.insert_right(Node(3))
n2.insert_left(Node(4))
assert t == w
assert t != Node(10)
assert t.preorder() == [1, 2, 4, 3]
assert t.inorder() == [4, 2, 1, 3]
assert t.postorder() == [4, 2, 3, 1]
# preorder: property: similar to dfs
# preorder: property: v[0] is the root
# inorder: property: left of element implies in left subtree
# inorder: property: right of element implies in right subtree
n = Node(3)
assert n.root() == n
assert n.root() is n
assert Node(3).root() == Node(3)
assert t.leaves() == [4, 3]
t.map(lambda x: x * 2)
assert t.leaves() == [8, 6]
w.clear()
assert len(w) == 0
assert not n2