C++ from recursive O(nlogn) to non-recursive O(n), all based on prefix sum


  • 0
    A

    I don't really like recursive functions, but still, here it is. O(nlogn)

    126 / 126 test cases passed
    Status: Accepted
    Runtime: 39 ms

    class Solution {
    public:
        void recur(vector<int>& v, int sum, TreeNode *root, uint& cnt) {
            if (!root) return;
            uint res = v.back() + root->val;
            for (int n : v) {
                if (res - n == sum) cnt++;
            }
            v.push_back(res); 
            recur(v, sum, root->left, cnt);
            recur(v, sum, root->right, cnt);
            v.pop_back();
        }
        int pathSum(TreeNode* root, int sum) {
            uint cnt = 0;
            vector<int> prefix_sum(1, 0);
            recur(prefix_sum, sum, root, cnt);
            return cnt;
        }
    };
    

    The non-recursive O(nlogn) solution.

    126 / 126 test cases passed
    Status: Accepted
    Runtime: 36 ms

    class Solution {
    public:
        int pathSum(TreeNode* root, int sum) {
            if (!root) return 0;
            uint cnt = 0;
            vector<int> prefix_sum(1, 0);
            stack<TreeNode*> stk;
            stk.push(root);
            while (!stk.empty()) {
                TreeNode* node = stk.top();
                if (!node) {
                    prefix_sum.pop_back();
                    stk.pop();
                } else {
                    uint res = prefix_sum.back() + node->val;
                    for (int n : prefix_sum) {
                        if (res - n == sum) cnt++;
                    }
                    prefix_sum.push_back(res);
                    stk.top() = NULL;
                    if (node->left) stk.push(node->left);
                    if (node->right) stk.push(node->right);
                }
            }
            return cnt;
        }
    };
    

    The optimized non-recursive O(n) solution.

    126 / 126 test cases passed
    Status: Accepted
    Runtime: 19 ms

    class Solution {
    public:
        int pathSum(TreeNode* root, int sum) {
            if (!root) return 0;
            uint cnt = 0;
            unordered_map<int, uint> prefix_sum;
            prefix_sum[0]++;
            stack<pair<TreeNode*, uint>> stk;
            stk.emplace(root, root->val);
            while (!stk.empty()) {
                auto& p = stk.top();
                if (!p.first) {
                    prefix_sum[p.second]--;
                    stk.pop();
                } else {
                    auto it = prefix_sum.find(p.second - sum);
                    if (it != prefix_sum.end()) cnt += it->second;
                    prefix_sum[p.second]++;
                    if (p.first->left) stk.emplace(p.first->left, p.second + p.first->left->val);
                    if (p.first->right) stk.emplace(p.first->right, p.second + p.first->right->val);
                    p.first = NULL;
                }
            }
            return cnt;
        }
    };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.