33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
from dataclasses import dataclass
|
|
from enum import Enum, auto
|
|
from typing import Iterator, Optional
|
|
|
|
|
|
@dataclass
|
|
class Node[T]:
|
|
value: T
|
|
left: Optional['Node[T]']
|
|
right: Optional['Node[T]']
|
|
|
|
|
|
type BinaryTree[T] = Optional[Node[T]]
|
|
|
|
class TraversalType(Enum):
|
|
INORDER, POSTORDER, PREORDER = auto(), auto(), auto()
|
|
|
|
def traverse[T](tree: BinaryTree[T], order: TraversalType = TraversalType.INORDER) -> Iterator[T]:
|
|
match (tree, order):
|
|
case (Node(value, left, right), TraversalType.INORDER):
|
|
yield from traverse(left, order)
|
|
yield value
|
|
yield from traverse(right, order)
|
|
case (Node(value, left, right), TraversalType.POSTORDER):
|
|
yield from traverse(left, order)
|
|
yield from traverse(right, order)
|
|
yield value
|
|
case (Node(value, left, right), TraversalType.PREORDER):
|
|
yield value
|
|
yield from traverse(left, order)
|
|
yield from traverse(right, order)
|
|
case _:
|
|
return |