Could you do something like this?
class TreeNode:
def __init__(self, x):
self.val = x
self.left = None
self.right = None
def display(self):
lines, *_ = self._display_aux()
for line in lines:
print(line)
def _display_aux(self):
"""Returns list of strings, width, height, and horizontal coordinate of the root."""
# No child.
if self.right is None and self.left is None:
line = '%s' % self.val
width = len(line)
height = 1
middle = width // 2
return [line], width, height, middle
# Only left child.
if self.right is None:
lines, n, p, x = self.left._display_aux()
s = '%s' % self.val
u = len(s)
first_line = (x + 1) * ' ' + (n - x - 1) * '_' + s
second_line = x * ' ' + '/' + (n - x - 1 + u) * ' '
shifted_lines = [line + u * ' ' for line in lines]
return [first_line, second_line] + shifted_lines, n + u, p + 2, n + u // 2
# Only right child.
if self.left is None:
lines, n, p, x = self.right._display_aux()
s = '%s' % self.val
u = len(s)
first_line = s + x * '_' + (n - x) * ' '
second_line = (u + x) * ' ' + '\\' + (n - x - 1) * ' '
shifted_lines = [u * ' ' + line for line in lines]
return [first_line, second_line] + shifted_lines, n + u, p + 2, u // 2
# Two children.
left, n, p, x = self.left._display_aux()
right, m, q, y = self.right._display_aux()
s = '%s' % self.val
u = len(s)
first_line = (x + 1) * ' ' + (n - x - 1) * \
'_' + s + y * '_' + (m - y) * ' '
second_line = x * ' ' + '/' + \
(n - x - 1 + u + y) * ' ' + '\\' + (m - y - 1) * ' '
if p < q:
left += [n * ' '] * (q - p)
elif q < p:
right += [m * ' '] * (p - q)
zipped_lines = zip(left, right)
lines = [first_line, second_line] + \
[a + u * ' ' + b for a, b in zipped_lines]
return lines, n + m + u, max(p, q) + 2, n + u // 2
def array_to_tree(arr):
# Handle case of empty array
if not arr:
return None
# Initialize root node and a queue
root = TreeNode(arr[0] if arr[0] != 'x' else None)
queue = [root]
# Initialize an index to traverse the array
index = 1
while queue and index < len(arr):
node = queue.pop(0)
# Handle left child
if arr[index] != 'x':
"""
If the value is 'x', then the left child is None.
We know this is the left node because we are following a level-order traversal.
"""
node.left = TreeNode(arr[index])
"""
Add the left child to the queue so we can handle its children later.
"""
queue.append(node.left)
"""
We must increment the index regardless of whether we handled the left child otherwise we will get stuck in an infinite loop.
"""
index += 1
# Handle right child if there are still nodes left
if index < len(arr) and arr[index] != 'x':
"""
If the index is less than the length of the array and the value is 'x', then the right child is None according to level-order traversal.
"""
node.right = TreeNode(arr[index])
"""
Just as we did with the left child, we add the right child to the queue so we can handle its children later.
"""
queue.append(node.right)
"""
We must increment the index regardless of whether we handled the right child otherwise we will get stuck in an infinite loop.
"""
index += 1
"""
Note: We intentionally increment the index twice since we are handling two children at a time.
"""
return root
So this:
if __name__ == "__main__":
arr = ["6", "x", "8", "4", "x"]
root = array_to_tree(arr)
# xref: https://stackoverflow.com/questions/34012886/print-binary-tree-level-by-level-in-python
root.display()
Would output this:
6_
\
8
/
4