GATE-CS-2006 Question 50

题目

给定一棵二叉树 $T$ 和一个整数 $k$。编写一个函数，找出从根节点到叶节点的所有路径中，和为 $k$ 的路径的个数。路径不需要从根节点开始，也不需要在叶节点结束，但必须是从父节点到子节点的的连续路径。

函数的原型为：

int count_paths_with_sum(Node *root, int k)

假设二叉树中最多有 $n$ 个节点，设计一个时间复杂度为 $O(n)$ 的算法。

思路

深度优先搜索：

从根开始，记录下来从根节点到当前节点的路径上的所有节点的和，然后从当前节点开始，递归查找其左子树和右子树是否有符合条件的路径。

具体实现细节详见代码解释。

代码

class Node:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def count_paths_with_sum(root: Node, k: int) -> int:
    """
    返回从根节点到叶节点的所有路径中，和为 k 的路径的个数。

    Args:
        root: 二叉树的根节点
        k: 要求的路径和

    Returns:
        int: 符合条件的路径个数，时间复杂度为 O(n)

    """
    count = 0  # 符合条件的路径个数
    path_count = {0: 1}  # 记录路径和出现的次数

    def dfs(node: Node, target: int, path_sum: int) -> None:
        nonlocal count, path_count

        if node is None:
            return

        path_sum += node.val
        # 如果当前路径和减去目标值所得的差值在之前的路径上出现过，
        # 则说明从该路径上的某一个节点到当前节点的路径上的节点和为目标值。
        count += path_count.get(path_sum - target, 0)
        # 记录当前路径和出现的次数
        path_count[path_sum] = path_count.get(path_sum, 0) + 1

        dfs(node.left, target, path_sum)
        dfs(node.right, target, path_sum)

        # 返回上一级路径前，需要将当前路径和出现的次数 - 1
        # 因为如果继续往上走，当前节点已经不在路径上了。
        path_count[path_sum] -= 1

    dfs(root, k, 0)
    return count

参考文献

Cracking the Coding Interview, 6th Edition