Johnny Metz
Johnny Metz

Reputation: 5965

Build a binary tree from a Python list

Let's say I have the following binary tree:

enter image description here

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]

enter image description here

Upvotes: 6

Views: 7549

Answers (4)

trincot
trincot

Reputation: 350137

I use the following function when solving LeetCode problems. It has these characteristics:

  • It really uses LeetCode's input format, so it does not follow the heap-organisation, but skips entries when the corresponding parent is already None -- contrary to what the heap-format would do.
  • It uses a queue so to insert nodes in the correct (breadth-first) order
  • It uses an iterator over the input list

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

hiro protagonist
hiro protagonist

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

Woody1193
Woody1193

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

Hasnain Khan
Hasnain Khan

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

Related Questions