Program to find k-length paths on a binary tree in Python

Finding k-length paths in a binary tree is a classic dynamic programming problem on trees. We need to count all unique paths of exactly k nodes, where paths can traverse both upward (child to parent) and downward (parent to child).

Problem Understanding

Given a binary tree with unique values and a number k, we count paths of length k. Two paths are different if they contain different nodes or visit nodes in different order.

Algorithm Approach

We use depth-first search with dynamic programming. For each node, we track how many paths of each length end at that node. Then we combine paths from left and right subtrees to form longer paths.

Implementation

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

class Solution:
    def solve(self, root, K):
        def dfs(node):
            if not node:
                return [1] + [0] * (K-1)
            
            left = dfs(node.left)
            right = dfs(node.right)
            
            # Count paths that pass through current node
            for i in range(K):
                self.ans += left[i] * right[K - 1 - i]
            
            # Build result array for current node
            res = [0] * K
            res[0] = res[1] = 1  # Node itself and single edge paths
            
            for i in range(1, K - 1):
                res[i + 1] += left[i]   # Extend left paths
                res[i + 1] += right[i]  # Extend right paths
            
            return res
        
        self.ans = 0
        dfs(root)
        return self.ans

# Create test tree
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)

ob = Solution()
result = ob.solve(root, 3)
print(result)

4

How It Works

The algorithm maintains an array for each node where res[i] represents the number of paths of length i ending at that node. For each node, we:

• Get path counts from left and right subtrees
• Combine paths from both subtrees that together form k-length paths through current node
• Build the result array by extending existing paths with the current node

Time and Space Complexity

Conclusion

This dynamic programming solution efficiently counts k-length paths by combining path information from subtrees. The key insight is tracking paths of different lengths ending at each node and combining them appropriately.