| 1 | | from json import dumps |
| 2 | | |
| 3 | | |
| 4 | 1 | class Tree: |
| 5 | 2 | def __init__(self, label, children=None): |
| 6 | 3 | self.label = label |
| 7 | 4 | self.children = children if children is not None else [] |
| 8 | | |
| 5 | |
| 9 | 6 | def __dict__(self): |
| 10 | | return {self.label: [c.__dict__() for c in sorted(self.children)]} |
| 11 | | |
| 12 | | def __str__(self, indent=None): |
| 13 | | return dumps(self.__dict__(), indent=indent) |
| 14 | | |
| 15 | | def __lt__(self, other): |
| 16 | | return self.label < other.label |
| 17 | | |
| 7 | return {self.label: [child.__dict__() for child in self.children]} |
| 8 | |
| 18 | 9 | def __eq__(self, other): |
| 10 | if not isinstance(other, Tree): |
| 11 | return False |
| 19 | 12 | return self.__dict__() == other.__dict__() |
| 20 | 13 | |
| 21 | 14 | def from_pov(self, from_node): |
| 22 | | pass |
| 23 | | |
| 15 | # Find the target node and build parent map |
| 16 | parent_map = {} |
| 17 | target_node = self._find_node_and_build_parent_map(from_node, parent_map, None) |
| 18 | |
| 19 | # Edge Case: Node not found |
| 20 | if target_node is None: |
| 21 | raise ValueError("Tree could not be reoriented") |
| 22 | |
| 23 | # Rebuild tree from the perspective of the target node |
| 24 | return self._rebuild_tree_from_node(target_node, parent_map) |
| 25 | |
| 24 | 26 | def path_to(self, from_node, to_node): |
| 25 | | pass |
| 27 | # Edge Case: Same start and end node |
| 28 | if from_node == to_node: |
| 29 | return [from_node] |
| 30 | |
| 31 | # Get tree from the perspective of from_node |
| 32 | try: |
| 33 | tree_from_pov = self.from_pov(from_node) |
| 34 | except ValueError: |
| 35 | # Edge Case: from_node doesn't exist |
| 36 | raise ValueError("Tree could not be reoriented") |
| 37 | |
| 38 | # Find path to to_node in the reoriented tree |
| 39 | path = tree_from_pov._find_path_to_node(to_node) |
| 40 | |
| 41 | # Edge Case: to_node doesn't exist in the tree |
| 42 | if path is None: |
| 43 | raise ValueError("No path found") |
| 44 | |
| 45 | return path |
| 46 | |
| 47 | def _find_node_and_build_parent_map(self, target_label, parent_map, parent): |
| 48 | if self.label == target_label: |
| 49 | return self |
| 50 | |
| 51 | for child in self.children: |
| 52 | parent_map[child.label] = self |
| 53 | found = child._find_node_and_build_parent_map(target_label, parent_map, self) |
| 54 | if found: |
| 55 | return found |
| 56 | |
| 57 | return None |
| 58 | |
| 59 | def _rebuild_tree_from_node(self, node, parent_map): |
| 60 | new_children = [] |
| 61 | |
| 62 | # Add original children (except the parent in the new orientation) |
| 63 | for child in node.children: |
| 64 | if parent_map.get(node.label) != child: # Don't add the parent as a child |
| 65 | new_children.append(self._rebuild_tree_from_node(child, parent_map)) |
| 66 | |
| 67 | # Add the parent as a child if it exists |
| 68 | parent = parent_map.get(node.label) |
| 69 | if parent: |
| 70 | # Create a new parent node without the current node as its child |
| 71 | parent_children = [c for c in parent.children if c.label != node.label] |
| 72 | new_parent = Tree(parent.label, parent_children) |
| 73 | new_children.append(self._rebuild_tree_from_node(new_parent, parent_map)) |
| 74 | |
| 75 | return Tree(node.label, new_children) |
| 76 | |
| 77 | def _find_path_to_node(self, target_label): |
| 78 | if self.label == target_label: |
| 79 | return [self.label] |
| 80 | |
| 81 | for child in self.children: |
| 82 | path = child._find_path_to_node(target_label) |
| 83 | if path: |
| 84 | return [self.label] + path |
| 85 | |
| 86 | return None |