Write a python script to experiment with depth first search.
This commit is contained in:
Tom Alexander
156
python_experiments/depth_first_search/depth_first_search.py
Executable file
156
python_experiments/depth_first_search/depth_first_search.py
Executable file
@@ -0,0 +1,156 @@
|
||||
#!/usr/bin/env python
|
||||
#
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Iterable
|
||||
|
||||
|
||||
def main():
|
||||
activity_tree = ActivityTree()
|
||||
activity_tree.activities[0] = Activity(id=0, parent=0, text="*")
|
||||
activity_tree.activities[1] = Activity(id=1, parent=0, text="A")
|
||||
activity_tree.activities[2] = Activity(id=2, parent=1, text="B")
|
||||
activity_tree.activities[3] = Activity(id=3, parent=2, text="C")
|
||||
activity_tree.activities[4] = Activity(id=4, parent=2, text="D")
|
||||
activity_tree.activities[5] = Activity(id=5, parent=4, text="E")
|
||||
activity_tree.activities[6] = Activity(id=6, parent=4, text="F")
|
||||
activity_tree.activities[7] = Activity(id=7, parent=1, text="G")
|
||||
activity_tree.activities[8] = Activity(id=8, parent=1, text="H")
|
||||
activity_tree.activities[9] = Activity(id=9, parent=8, text="I")
|
||||
activity_tree.activities[10] = Activity(id=10, parent=9, text="J")
|
||||
|
||||
match_iter = MatchIter(activity_tree, lambda act: act.id in (2, 5, 10))
|
||||
|
||||
while (depth_and_node := match_iter.get_next()) is not None:
|
||||
depth, matched, node = depth_and_node
|
||||
print("\t".join((str(depth), str(matched), node.text)))
|
||||
|
||||
# while True:
|
||||
# print("\n\n\nvvvvvvvvvv\n\n")
|
||||
# node = match_iter.get_next()
|
||||
# if node is None:
|
||||
# break
|
||||
# print(node.text)
|
||||
# print("\n\n")
|
||||
# match_iter.print_stack()
|
||||
# print("\n\n\n^^^^^^^^^^\n\n")
|
||||
|
||||
|
||||
class ActivityTree:
|
||||
activities: dict[int, Activity]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.activities = {}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Activity:
|
||||
id: int
|
||||
parent: int
|
||||
text: str
|
||||
|
||||
|
||||
class MatchIter:
|
||||
activity_tree: ActivityTree
|
||||
stack: deque[StackEntry]
|
||||
is_match: Callable[[Activity], bool]
|
||||
|
||||
def __init__(
|
||||
self, activity_tree: ActivityTree, is_match: Callable[[Activity], bool]
|
||||
) -> None:
|
||||
self.activity_tree = activity_tree
|
||||
self.stack = deque([StackEntry(0)])
|
||||
self.is_match = is_match
|
||||
|
||||
def get_next(self) -> tuple[int, bool, Activity] | None:
|
||||
if (next_announce := self.get_next_announce()) is not None:
|
||||
next_announce_depth, next_announce_entry = next_announce
|
||||
next_announce_entry.announced = True
|
||||
return (
|
||||
next_announce_depth,
|
||||
next_announce_entry.matched,
|
||||
self.activity_tree.activities[next_announce_entry.node_id],
|
||||
)
|
||||
|
||||
while len(self.stack) > 0:
|
||||
stack_entry = self.stack[-1]
|
||||
node = self.activity_tree.activities[stack_entry.node_id]
|
||||
|
||||
if stack_entry.visited:
|
||||
self.stack.pop()
|
||||
continue
|
||||
# print(f"Checking {stack_entry.text(self.activity_tree)}")
|
||||
|
||||
if self.is_match(node):
|
||||
stack_entry.should_announce = True
|
||||
stack_entry.matched = True
|
||||
for parent in self.walk_up_stack(stack_entry):
|
||||
if parent.node_id != 0 and not parent.should_announce:
|
||||
parent.should_announce = True
|
||||
|
||||
stack_entry.visited = True
|
||||
|
||||
children = [
|
||||
c for c in self.activity_tree.activities.values() if c.parent == node.id
|
||||
]
|
||||
# for c in children:
|
||||
for c in reversed(children):
|
||||
if c.id != 0:
|
||||
self.stack.append(StackEntry(c.id))
|
||||
|
||||
if (
|
||||
stack_entry.should_announce
|
||||
and (next_announce := self.get_next_announce()) is not None
|
||||
):
|
||||
next_announce_depth, next_announce_entry = next_announce
|
||||
next_announce_entry.announced = True
|
||||
return (
|
||||
next_announce_depth,
|
||||
next_announce_entry.matched,
|
||||
self.activity_tree.activities[next_announce_entry.node_id],
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
def print_stack(self) -> None:
|
||||
for stack_entry in self.stack:
|
||||
print(stack_entry.text(self.activity_tree))
|
||||
|
||||
def get_next_announce(self) -> tuple[int, StackEntry] | None:
|
||||
for entry in self.stack:
|
||||
if entry.should_announce and not entry.announced:
|
||||
depth = len(list(self.walk_up_stack(entry)))
|
||||
|
||||
return (depth, entry)
|
||||
return None
|
||||
|
||||
def walk_up_stack(self, base: StackEntry) -> Iterable[StackEntry]:
|
||||
found_start = False
|
||||
for i in range(len(self.stack)):
|
||||
idx = -1 - i
|
||||
stack_entry = self.stack[idx]
|
||||
if not found_start:
|
||||
if stack_entry is base:
|
||||
found_start = True
|
||||
continue
|
||||
if stack_entry.visited:
|
||||
yield stack_entry
|
||||
|
||||
|
||||
@dataclass
|
||||
class StackEntry:
|
||||
node_id: int
|
||||
announced: bool = False
|
||||
should_announce: bool = False
|
||||
visited: bool = False
|
||||
matched: bool = False
|
||||
|
||||
def text(self, activity_tree: ActivityTree) -> str:
|
||||
node = activity_tree.activities[self.node_id]
|
||||
return f"{node.text}\t{"annced" if self.announced else "-"}\t{"should" if self.should_announce else "-"}\t{"visited" if self.visited else "-"}\t{"matched" if self.matched else "-"}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
106
python_experiments/draw_tree/draw_tree.py
Normal file
106
python_experiments/draw_tree/draw_tree.py
Normal file
@@ -0,0 +1,106 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
def main():
|
||||
activity_tree = ActivityTree()
|
||||
activity_tree.activities[1] = Activity(id=1, parent=0, text="foo")
|
||||
activity_tree.activities[2] = Activity(id=2, parent=0, text="bar")
|
||||
activity_tree.activities[3] = Activity(id=3, parent=0, text="baz")
|
||||
activity_tree.activities[4] = Activity(id=4, parent=2, text="lorem")
|
||||
# activity_tree.activities[5] = Activity(id=5, parent=2, text="ipsum")
|
||||
# activity_tree.activities[6] = Activity(id=6, parent=2, text="dolar")
|
||||
activity_tree.activities[7] = Activity(id=7, parent=4, text="north")
|
||||
activity_tree.activities[8] = Activity(id=8, parent=4, text="east")
|
||||
activity_tree.activities[9] = Activity(id=9, parent=4, text="south")
|
||||
activity_tree.activities[10] = Activity(id=10, parent=4, text="west")
|
||||
|
||||
# For each line we need to know:
|
||||
# - the depth
|
||||
# - whether there are later siblings
|
||||
# - whether there are earlier siblings
|
||||
#
|
||||
# Draw order:
|
||||
# - Find all leaves
|
||||
# - Start with last child of root
|
||||
# - If it has children, keep going down until you hit leaves
|
||||
|
||||
for entry in activity_tree.get_draw_order():
|
||||
text = activity_tree.activities[entry.id].text
|
||||
leading_bars = "".join(" " if depth else " " for depth in entry.depth)
|
||||
branch = "" if entry.has_later_siblings else ""
|
||||
print(f"{leading_bars}{branch} {text}")
|
||||
|
||||
|
||||
class ActivityTree:
|
||||
activities: dict[int, Activity]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.activities = {}
|
||||
|
||||
def get_draw_order(self) -> list[StackEntry]:
|
||||
draw_order = []
|
||||
stack = [
|
||||
StackEntry(
|
||||
id=child_id,
|
||||
depth=[],
|
||||
has_later_siblings=True,
|
||||
have_pulled_children=False,
|
||||
)
|
||||
for child_id in self.get_children_in_order(0)
|
||||
]
|
||||
stack[-1].has_later_siblings = False
|
||||
while len(stack) > 0:
|
||||
current_entry = stack.pop()
|
||||
if current_entry.have_pulled_children:
|
||||
if current_entry.id != 0:
|
||||
draw_order.append(current_entry)
|
||||
continue
|
||||
|
||||
current_entry.have_pulled_children = True
|
||||
stack.append(current_entry)
|
||||
all_children = self.get_children_in_order(current_entry.id)
|
||||
for child_id in all_children:
|
||||
stack.append(
|
||||
StackEntry(
|
||||
id=child_id,
|
||||
depth=current_entry.depth + [current_entry.has_later_siblings],
|
||||
has_later_siblings=True,
|
||||
have_pulled_children=False,
|
||||
)
|
||||
)
|
||||
if len(all_children) > 0:
|
||||
stack[-1].has_later_siblings = False
|
||||
return draw_order
|
||||
|
||||
def get_children_in_order(self, parent_id: int) -> list[int]:
|
||||
return sorted(
|
||||
child_id
|
||||
for child_id, child in self.activities.items()
|
||||
if child.parent == parent_id
|
||||
)
|
||||
|
||||
def has_children(self, parent_id: int) -> bool:
|
||||
return any(
|
||||
child.parent == parent_id for _child_id, child in self.activities.items()
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class StackEntry:
|
||||
id: int
|
||||
depth: list[bool]
|
||||
has_later_siblings: bool
|
||||
have_pulled_children: bool # Only used when building the draw order
|
||||
|
||||
|
||||
@dataclass
|
||||
class Activity:
|
||||
id: int
|
||||
parent: int
|
||||
text: str
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user