Path to each Leaf of a Binary Tree

Viewed 1256

The function above AllPaths() appends an array containing the path to each leaf of the binary tree to the global array res.

The code works just fine, but I want to remove the global variable res and make the function return an array instead. How can I do that?

class Node:
    def __init__(self, value, left=None, right=None) -> None:
        self.value = value
        self.left  = left
        self.right = right

res = []
def allPaths(node, arr=[]):
    if node:
        tmp = [*arr, node.value]
        if not node.left and not node.right: # Leaf
            res.append(tmp)
        allPaths(node.left, tmp)
        allPaths(node.right, tmp)


root             = Node(1)
root.left        = Node(2);
root.left.left   = Node(4);
root.left.right  = Node(5);
root.right       = Node(3);
root.right.right = Node(6);
"""
          1         <-- root
        /   \
       2     3  
     /   \    \
    4     5    6    <-- leaves
"""
allPaths(root)
print(res)
# Output : [[1, 2, 4], [1, 2, 5], [1, 3, 6]]
4 Answers

A simple way that allow you to avoid the inner lists and global list altogether is to make a generator that yields the values as they come. Then you can just pass this to list to make the final outcome:

class Node:
    def __init__(self, value, left=None, right=None) -> None:
        self.value = value
        self.left  = left
        self.right = right

def allPaths(node):  
    if node:
        if not node.left and not node.right: # Leaf
            yield [node.value]
        else:
            yield from ([node.value] + arr for arr in allPaths(node.left))
            yield from ([node.value] + arr for arr in allPaths(node.right))
              
root             = Node(1)
root.left        = Node(2);
root.left.left   = Node(4);
root.left.right  = Node(5);
root.right       = Node(3);
root.right.right = Node(6);
        
g = allPaths(root)
list(g)

# [[1, 2, 4], [1, 2, 5], [1, 3, 6]]

One method is to do it by backtracking:

def allPaths(node, partial_res, res):
    if not node: 
        return
    if not node.left and not node.right:
        res.append(partial_res[:] + [node.value])
        return    
    partial_res.append(node.value)
    allPaths(node.left, partial_res, res)
    allPaths(node.right, partial_res, res)
    partial_res.pop()

res = []
allPaths(root, [], res)
print(res)

You could pass down the current path in the recursion:

def allPaths(node,path=[]):
    if not node: return            # no node, do nothing
    fullPath = path + [node.value]
    if node.left or node.right:    # node is not a leaf, recurse down      
        yield from allPaths(node.left, fullPath)    # left leaves if any
        yield from allPaths(node.right, fullPath)   # right leaves if any
    else:
        yield fullPath             # leaf node, return final path

I offer another option.

def allPaths(root, path=[]):
    tmp = []
    if root.left:
        tmp.extend(allPaths(root.left, path + [root.value]))
    if root.right:
        tmp.extend(allPaths(root.right, path + [root.value]))
    if not root.left and not root.right:
        tmp.append(path + [root.value])
    return tmp


tree = allPaths(root)
print(tree)

The output is:

[[1, 2, 4], [1, 2, 5], [1, 3, 6]]
Related