Skip to content

Commit ab8cc89

Browse files
Fix tree visualization and test all helper functions
Co-authored-by: jeremymanning <9030494+jeremymanning@users.noreply.github.com>
1 parent 73647b4 commit ab8cc89

File tree

1 file changed

+30
-12
lines changed

1 file changed

+30
-12
lines changed

helpers/tree_utils.py

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -96,14 +96,15 @@ def tree_to_list(root: Optional[TreeNode]) -> List[Optional[int]]:
9696
return result
9797

9898

99-
def visualize_tree(root: Optional[TreeNode], level: int = 0, prefix: str = "Root: ") -> str:
99+
def visualize_tree(root: Optional[TreeNode], prefix: str = "", is_tail: bool = True, is_root: bool = True) -> str:
100100
"""
101101
Create a visual string representation of a binary tree.
102102
103103
Args:
104104
root: Root node of the binary tree
105-
level: Current depth level (used for recursion, default 0)
106-
prefix: Prefix string for the current node (default "Root: ")
105+
prefix: Prefix string for the current node (used internally)
106+
is_tail: Whether this is the last child (used internally)
107+
is_root: Whether this is the root node (default True)
107108
108109
Returns:
109110
String representation of the tree structure
@@ -121,17 +122,34 @@ def visualize_tree(root: Optional[TreeNode], level: int = 0, prefix: str = "Root
121122
return ""
122123

123124
lines = []
124-
lines.append(prefix + str(root.val))
125125

126-
if root.left or root.right:
127-
if root.left:
128-
extension = "│ " if root.right else " "
129-
lines.append(visualize_tree(root.left, level + 1,
130-
("│ " * level) + "├─ L: ").rstrip())
126+
# Root node or child node label
127+
if is_root:
128+
lines.append("Root: " + str(root.val))
129+
new_prefix = ""
130+
else:
131+
lines.append(prefix + str(root.val))
132+
new_prefix = prefix + (" " if is_tail else "│ ")
133+
134+
# Process children
135+
children = []
136+
if root.left:
137+
children.append(('L', root.left))
138+
if root.right:
139+
children.append(('R', root.right))
140+
141+
for i, (label, child) in enumerate(children):
142+
is_last = (i == len(children) - 1)
143+
connector = "└─ " if is_last else "├─ "
144+
child_lines = visualize_tree(child, new_prefix, is_last, False)
131145

132-
if root.right:
133-
lines.append(visualize_tree(root.right, level + 1,
134-
("│ " * level) + "└─ R: ").rstrip())
146+
if child_lines:
147+
# Add the label to the first line
148+
first_line = new_prefix + connector + label + ": " + str(child.val)
149+
remaining_lines = child_lines.split('\n')[1:] if '\n' in child_lines else []
150+
151+
lines.append(first_line)
152+
lines.extend(remaining_lines)
135153

136154
return "\n".join(lines)
137155

0 commit comments

Comments
 (0)