107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
|
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()
|