Reputation: 5965
Let's say I have the following binary tree:
and the following TreeNode
definition:
class TreeNode:
def __init__(self, x, left: = None, right: = None):
self.val = x
self.left = left
self.right = right
To construct this tree, I currently need to write the following code:
def build_tree() -> TreeNode:
node0 = TreeNode(2, left=TreeNode(4), right=TreeNode(5))
node1 = TreeNode(3, right=TreeNode(7))
root = TreeNode(1, left=node0, right=node1)
return root
This is very inefficient for large trees. Leetcode gives me trees as a Python list. For the example above, it reads:
[1, 2, 3, 4, 5, None, 7]
I would love to refactor build_tree
so it takes a list, builds the tree and returns the root node:
def build_tree(values: list) -> ListNode:
...
Any ideas on how to create this function?
Here's a more complicated example:
[4, 6, 2, None, 5, 1, 9, 2, 1, 3, 4, None, 7]
Upvotes: 6
Views: 7549
Reputation: 350137
I use the following function when solving LeetCode problems. It has these characteristics:
None
-- contrary to what the heap-format would do.Code:
def to_binary_tree(items):
if not items:
return None
it = iter(items)
root = TreeNode(next(it))
q = [root]
for node in q:
val = next(it, None)
if val is not None:
node.left = TreeNode(val)
q.append(node.left)
val = next(it, None)
if val is not None:
node.right = TreeNode(val)
q.append(node.right)
return root
Upvotes: 2
Reputation: 46849
assuming your list is already in the order described by https://docs.python.org/3/library/heapq.html?highlight=heapq#theory you could do this (if it is not, use heapq.heapify
):
def from_list(elements):
root_node = TreeNode(x=elements[0])
nodes = [root_node]
for i, x in enumerate(elements[1:]):
if x is None:
continue
parent_node = nodes[i // 2]
is_left = (i % 2 == 0)
node = TreeNode(x=x)
if is_left:
parent_node.left = node
else:
parent_node.right = node
nodes.append(node)
return root_node
i store all the nodes in a list called nodes
. iterating over the elements
i then know how to find the parent node in the nodes
list.
using the recursive print function
def print_recursive(tree, level=0, prefix="root"):
print(f"{level * ' '}{prefix:5s}: val={tree.val}")
if tree.left:
print_recursive(tree.left, level + 1, "left")
if tree.right:
print_recursive(tree.right, level + 1, "right")
and running it for your second example
tree = from_list(elements=[4, 6, 2, None, 5, 1, 9, 2, 1, 3, 4, None, 7])
print_recursive(tree)
i get:
root : val=4
left : val=6
right: val=5
left : val=2
right: val=1
right: val=2
left : val=1
left : val=3
right: val=4
right: val=9
right: val=7
Upvotes: 2
Reputation: 7970
You could iterate over your elements and convert the list to a tree this way. This isn't very memory efficient (it will store a pointer to each node on the queue so you're looking at O(N) in terms of data) but it will be more efficient in terms of CPU, especially in Python.
def build_tree(data):
# Check if we have data; if we don't then return None
if not data:
return None
# Create our tree object from the root and then
# add it to a list of nodes that we can add data to
tree = TreeNode(data[0])
queue = [data[0]]
# Iterate over the rest of the nodes and add
# them to the tree
index = 0
for i in range(1, len(data) - 1):
# If our node is None then we have a leaf so
# there's no work to do; otherwise, create the
# node and add it to the proper spot
if data[i] != None:
# Create the node from the data
node = TreeNode(data[i])
# If the index we're looking at is even
# then we have a left branch; otherwise we
# have a right branch so add the node to the
# appropriate side of the tree
if i % 2 == 0:
queue[index].left = node
else:
queue[index].right = node)
# Add the node to our queue so it can
# be referenced later
queue.append(node)
# Finally, update our index if we've finished looking
# at this node. As each branch may only have two children
# we only need to do this if we're looking at odd-indexed
# entries because leaves will always be None (unless we've
# reached the last level of the tree)
index += i % 2
return tree
This algorithm works by building a queue from the data (which already had a heap structure anyway) while creating the nodes. Because the queue is iterated only when the current node is full and because it iterates left-to-right, the order of the tree will be preserved.
Upvotes: 0
Reputation: 373
You can refactor the code like this. It will take a list of numbers or string and will build a BST but for simple binary tree, only one if
condition will needs to be removed.
def add_child(self, val):
if val == self.val:
return
elif val < self.val:
if self.left:
self.left.add_child(val)
else:
self.left = BinarySearchTree(val)
else:
if self.right:
self.right.add_child(val)
else:
self.right = BinarySearchTree(val)
def build_tree(elements):
root = BinarySearchTree(elements[0])
for i in range(1, len(elements)):
root.add_child(elements[i])
return root
if __name__=="__main__":
numbers = [15, 10, 8, 12, 20, 16, 25]
root = build_tree(numbers)
Upvotes: 0