路径总和 III - 前缀和还能这样用!

309 阅读4分钟

小知识,大挑战!本文正在参与“程序员必备小知识”创作活动。

题目介绍

题目详情可看这里

给定一个二叉树的根节点 root ,和一个整数 targetSum ,求该二叉树里节点值之和等于 targetSum 的 路径 的数目。

路径 不需要从根节点开始,也不需要在叶子节点结束,但是路径方向必须是向下的(只能从父节点到子节点)。

示例:

pathsum3-1-tree.jpg

输入:root = [10,5,-3,3,2,null,11,3,-2,null,1], targetSum = 8
输出:3
解释:和等于 8 的路径有 3 条,如图所示。

题目解读

首先题目要求计算所有符合的路径,那必然是需要遍历整个树的。而要遍历树,最常见的方法就是 dfs 和 bfs 了,这题我就使用 dfs 来对树进行遍历。

其次路径并不需要从根节点开始,所以任意一个结点都可以作为路径的开头,理解到这里最直观的想法就是把每个结点都当作是一个根节点,再往下去遍历搜寻合适的路径。

朴素解法

那么首先我想到的方法(当然先是想到笨办法!)是在 dfs 的过程中去维护一个 list,这个 list 记录从根节点到某个结点所有结点的值。然后在每次遍历到某个结点的时候,从后往前的遍历整个 list 并计算路径的和(即相当于以当前结点为路径的尾所能得到的全部路径的和),这样的话就不用重复的去做 dfs 了。

当然还要注意回溯,在回溯的时候要把不应该存在于 list 中的结点从list中抹除。

具体来看看我的实现代码:

class Solution {
    public int pathSum(TreeNode root, int targetSum) {
        List<Integer> sum = new ArrayList<>();

        if (root == null) return 0;
        return dfs(root, sum, targetSum);
    }

    public int dfs(TreeNode curr, List<Integer> sum, int targetSum) {
        int ans = 0;

        if (curr.val == targetSum) ans++;
        for (int i = sum.size() - 1, tmp = curr.val; i >= 0; i--) {
            tmp += sum.get(i);
            if (tmp == targetSum) ans++;
        }

        sum.add(curr.val);
        if (curr.left != null) ans += dfs(curr.left, sum, targetSum);
        if (curr.right != null) ans += dfs(curr.right, sum, targetSum);
        sum.remove(sum.size() - 1);

        return ans;
    }
}

除此之外也可以使用前面想到的最直观的做法,也就是以每个结点为根节点(路径的开头)去寻找所有可能的路径。代码如下:

class Solution {
    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) return 0;
        return dfs(root, targetSum);
    }

    public int dfs(TreeNode curr, int targetSum) {
        int ans = ddfs(curr, curr.val, targetSum);
        if (curr.left != null) ans += dfs(curr.left, targetSum);
        if (curr.right != null) ans += dfs(curr.right, targetSum);

        return ans;
    }

    public int ddfs(TreeNode curr, int curSum, int targetSum) {
        int ans = 0;

        if (curSum == targetSum) ans += 1;

        if (curr.left != null) ans += ddfs(curr.left, curSum + curr.left.val, targetSum);
        if (curr.right != null) ans += ddfs(curr.right, curSum + curr.right.val, targetSum);

        return ans;
    }
}

前缀和的应用

上面的方法我们很容易想到,计算路径的过程中存在许多无用的重复计算。比如在计算 1 - 2 - 3 - 4 - 5 这条路径的时候,每经过一个结点都会再次重复计算前面所有结点的路径和。

从这我们又能想到,针对一个线性的数组来算路径和的时候我们通常都是这么做的?当然是用前缀和了!如 1 - 2 - 3 - 4 - 5 这条路径,我们把当前位置置为前面所有值与当前值的总和,就得到了前缀和 1 - 3 - 6 - 10 - 15。这样的话要想得到某个位置到某个位置之间的路径和时只需把对应位置的前缀和元素一减即可得到,如要得到位置2位置4的路径和,那就是 10 - 1 = 92 + 3 + 4

用在这题也是一样的,路径都是一条线,在遍历的时候存储根节点到当前节点的前缀和就可以了,完全可以把他当成是线性的。只是在回溯的时候要记得把应该去掉的前缀和删除。

在实现细节上,使用哈希表来记录前缀和,毕竟这个题的某个路径和可能不止存在1个。

具体代码如下:

class Solution {
    int t;

    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) return 0;
        Map<Integer, Integer> map = new HashMap<>();
        t = targetSum;

        map.put(0, 1);
        return dfs(root, root.val, map);
    }

    public int dfs(TreeNode curr, int sum, Map<Integer, Integer> map) {
        int ans = 0;
        if (map.containsKey(sum - t))  ans += map.get(sum - t);

        map.put(sum, map.getOrDefault(sum, 0) + 1);
        if (curr.left != null) ans += dfs(curr.left, sum + curr.left.val, map);
        if (curr.right != null) ans += dfs(curr.right, sum + curr.right.val, map);
        map.put(sum, map.getOrDefault(sum, 0) - 1);

        return ans;
    }
}

总结

从这个题可以看出,遇到问题先想想最朴素的做法,然后从这个做法的缺点上入手再去寻找改进的思路,就能得到比较好的解法了。

除此之外前缀和也是一个很好的想法,不仅可以应用在线性的数据结构上,在树上同样能得到很好的应用。