Write a python script to experiment with depth first search.

python_experiments/depth_first_search/depth_first_search.py

@@ -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()