【LeetCode】113. Path Sum II

Path Sum II

Given a binary tree and a sum, find all root-to-leaf paths where each path's sum equals the given sum.

For example:
Given the below binary tree and sum = 22,

              5
             / \
            4   8
           /   / \
          11  13  4
         /  \    / \
        7    2  5   1

return

[
   [5,4,11,2],
   [5,8,4,5]
]

 

这题是在Path Sum的基础上稍作修改。

与上题增加的动作是保存当前stack中路径的加和curSum。

当curSum等于sum时，将cur加入ret。

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<vector<int> > pathSum(TreeNode *root, int sum) {
        vector<vector<int> > ret;
        if(!root)
            return ret;
            
        stack<TreeNode*> stk;
        vector<int> cur;
        cur.push_back(root->val);
        int curSum = root->val;
        unordered_map<TreeNode*, bool> visited;
        stk.push(root);
        visited[root] = true;
        
        while(!stk.empty())
        {
            TreeNode* top = stk.top();
            if(!top->left && !top->right)
            {//leaf
                if(curSum == sum)
                    ret.push_back(cur);
            }
            
            if(top->left && visited[top->left] == false)
            {
                stk.push(top->left);
                visited[top->left] = true;
                cur.push_back(top->left->val);
                curSum += top->left->val;
                continue;
            }
            if(top->right && visited[top->right] == false)
            {
                stk.push(top->right);
                visited[top->right] = true;
                cur.push_back(top->right->val);
                curSum += top->right->val;
                continue;
            }
            
            stk.pop();
            curSum -= top->val;
            cur.pop_back();
        }
        return ret;
    }
};