Skip to content

树上前缀和

Prefix sum in a tree

用字典来记录当前节点到根节点路径上的前缀和们。

class Solution {
public:
    int pathSum(TreeNode* root, int targetSum) {
        // dfs
        int ans = 0;
        unordered_map<int, int> m;
        m[0] = 1;

        function<void(TreeNode*, int)> find = [&](TreeNode* r, int s) {
            if (r == nullptr) return;
            int ss = s + r->val;
            ans += m[ss - targetSum];
            m[ss]++;
            find(r->left, ss);
            find(r->right, ss);
            m[ss]--;
        };

        find(root, 0);
        return ans;
    }
};