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