Count of Range Sum | LeetCode

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.
Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.

Example:
Given nums = [-2, 5, -1], lower = -2, upper = 2,
Return 3.
The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

思路：对这种range sum的题，首先的思路就是计算range时，可以先用一个presum，然后用presum[j]-presum[i-1]来计算sum[i~j]的值。 所以这里我们一般在presum前补一个0，来使得sum[1~i] = presum[i] - presume[0]。然后本来我很自然的想到sort一下，但是sort之后会打乱原有的顺序，所以不可取。再进一步，在tag里面看到提示，binary search tree和原先做过count of smaller numbers after self, 就想到了维护一个二叉搜索树。

参考了(https://leetcode.com/discuss/79251/a-c-binary-search-tree-solution)

对于信插入的元素x，我们只要在树中找总共有多少presum满足下式

x - presume[i] >= lower && x -presume[i] <= upper 即可。

这里我们用了 x-presum[i] <= upper 的数量，减去 ，x-presum[i] < lower 的数量来求得。

代码如下：

class BTNode {
public:
    long val;
    int rightchild;
    BTNode* left;
    BTNode* right;
    BTNode(long a):val(a),rightchild(0),left(NULL),right(NULL){};
};

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        vector<long> presum(1,0);
        for (auto &i:nums) {
            presum.push_back(presum.back()+i);
        }
        BTNode root(0);
        int result = 0;
        for (int i = 1; i < presum.size(); i++) {
            long target = presum[i];
            int a = largerthan(&root,target-lower);
            int b = largerthan(&root,target-upper-1);
            result += b-a;
            insert(&root,target);
        }
        return result;
    }
    
    int largerthan(BTNode*root, long target) {
        if (!root) return 0;
        if (root->val == target) return root->rightchild;
        if (target < root->val) return root->rightchild+largerthan(root->left,target)+1;
        if (target > root->val) return largerthan(root->right,target);
        return 0;
    }
    
    void insert(BTNode* root,int target) {
        while (root) {
            if (target <= root->val) {
                if (root->left) {
                    root = root->left;
                } else {
                    root->left = new BTNode(target);
                    break;
                }
            } else {
                root->rightchild++;
                if (root->right) {
                    root = root->right;
                } else {
                    root->right = new BTNode(target);
                    break;
                }
            }
        }
        return;
    }
    
};