5 Efficient Ways to Find K Length Paths in a Binary Tree using Python

Rate this post

πŸ’‘ Problem Formulation: Given a binary tree and an integer k, the challenge is to find all the paths in the tree that consist of exactly k edges. A path in a binary tree is a sequence of nodes where each pair of adjacent nodes has an edge connecting them. The input consists of the root of the binary tree and the integer k. The desired output is a list of paths, with each path represented as a list of node values.

Method 1: Recursive Depth-First Search (DFS)

The DFS algorithm can be used to traverse the tree and keep track of the current path’s length. When the path’s length reaches k, the path can be added to the resulting list. Here’s a function find_k_length_paths(root, k), which returns all k-length paths starting at the root node.

Here’s an example:

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

def dfs(node, k, path, result):
    if not node: return
    if k == 0:
        result.append(path + [node.val])
    dfs(node.left, k-1, path + [node.val], result)
    dfs(node.right, k-1, path + [node.val], result)

def find_k_length_paths(root, k):
    result = []
    dfs(root, k, [], result)
    return result

# Example binary tree
root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)

print(find_k_length_paths(root, 2))

Output:

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

This implementation of DFS starts by calling dfs() on the root node with an empty path and the desired length k. The dfs() function explores each child node, decrementing k until it reaches zero, meaning a valid path has been found, which is then added to the results. This approach is simple but may include extra calculations for paths longer than k.

Method 2: Iterative DFS with Stack

This method emulates recursive DFS using an explicit stack. At each step, it maintains a tuple containing the current node, the remaining length of k to be traversed, and the current path. The benefit of this method is the control it gives to iterate through the tree without the extra memory overhead of system call stack frames.

Here’s an example:

def iterative_dfs(root, k):
    if root is None: return []
    stack, result = [(root, k, [])], []

    while stack:
        node, remaining_k, path = stack.pop()
        if remaining_k == 0:
            result.append(path + [node.val])
        if node.left:
            stack.append((node.left, remaining_k-1, path + [node.val]))
        if node.right:
            stack.append((node.right, remaining_k-1, path + [node.val]))

    return result

print(iterative_dfs(root, 2))

Output:

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

The iterative_dfs() function uses a stack to manage the traversal state instead of system stack frames, providing an alternative for systems with limited recursive depth. It works similarly to the recursive version but is manually controlling the stack of nodes to visit.

Method 3: Breadth-First Search (BFS) with Queue

BFS traverses the tree level by level and can also solve this problem. In this method, a queue maintains the current node alongside the current path and the remaining length of k. When k reaches zero, the corresponding path is added to the solutions.

Here’s an example:

from collections import deque

def bfs(root, k):
    if root is None: return []
    queue, result = deque([(root, k, [])]), []

    while queue:
        node, remaining_k, path = queue.popleft()
        if remaining_k == 0:
            result.append(path + [node.val])
        if node.left:
            queue.append((node.left, remaining_k-1, path + [node.val]))
        if node.right:
            queue.append((node.right, remaining_k-1, path + [node.val]))

    return result

print(bfs(root, 2))

Output:

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

The bfs() function uses a queue and processes each level of the tree before moving deeper. This can be more memory-efficient for wide trees, as it stops the traversal after the specified depth k is reached, unlike the DFS which may continue deeper in the tree.

Method 4: Backtracking

Backtracking is a refined DFS approach that reuses the path and adjusts it while traversing the tree to avoid unnecessary memory allocations for each path. This method is more memory efficient because there’s no need to create new path lists at every node.

Here’s an example:

def backtrack_paths(node, k, path, result):
    if not node: return
    path.append(node.val)
    if k == 0:
        result.append(list(path))
    else:
        backtrack_paths(node.left, k-1, path, result)
        backtrack_paths(node.right, k-1, path, result)
    path.pop()

def find_paths_backtracking(root, k):
    result, path = [], []
    backtrack_paths(root, k, path, result)
    return result

print(find_paths_backtracking(root, 2))

Output:

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

This backtrack_paths() function optimizes memory usage by mutating a single path list, pushing and popping nodes values as it traverses the tree. When a path of length k is reached, a copy of the path is stored in the results.

Bonus One-Liner Method 5: Using Python Generators

Python generators allow us to create concise and memory-efficient solutions. This method uses a generator to yield paths of length k. The main function find_k_length_paths_gen can be used to generate k-length paths on demand without storing all paths in memory at once.

Here’s an example:

def paths_gen(node, k, path):
    if not node: return
    if k == 0:
        yield path + [node.val]
    yield from paths_gen(node.left, k-1, path + [node.val])
    yield from paths_gen(node.right, k-1, path + [node.val])

def find_k_length_paths_gen(root, k):
    return list(paths_gen(root, k, []))

print(find_k_length_paths_gen(root, 2))

Output:

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

This approach utilizes Python’s yield statement and yield from expression to efficiently produce k-length paths as the recursion unfolds. The generator avoids holding all results in memory at once and can be iterated over to retrieve one path at a time.

Summary/Discussion

  • Method 1: Recursive DFS: Is intuitive and easy to implement. Recursion can lead to a stack overflow on very deep trees.
  • Method 2: Iterative DFS with Stack: Controls stack usage manually, avoiding recursion limits. Can be slightly more complex to reason about.
  • Method 3: BFS with Queue: Is memory-efficient for trees that aren’t deep but wide. Ensures early termination after reaching the desired depth.
  • Method 4: Backtracking: Saves memory by reusing the path list. Requires careful handling to avoid side effects.
  • Bonus One-Liner Method 5: Python Generators: Memory-efficient for generating and processing paths on the fly. Potentially complex for beginners to comprehend.