BST solution in 192ms


  • -1
    8
    class Node
    {
    public:
      long long num;
      int size;
      Node *ls,*rs;
      Node()
      {
          ls=NULL;
          rs=NULL;
          size=1;
      }
    };
    
    class Solution {
    public:
        int countRangeSum(vector<int>& nums, int lower, int upper) {
            init();
            int ans=0;
    		long long presum=0;
            
            insert(0,root);
            for(int i=0;i<nums.size();i++)
            {
                presum+=nums[i];
                long long MT=presum-upper;
                long long LT=presum-lower;
                ans+=queryLT(LT,root)-queryLT(MT-1,root);
                insert(presum,root);
            }
            return ans;
        }
    private:
        Node *root;
        void init()
        {
            root=NULL;
        }
        void insert(long long num,Node *&cur)
        {
            if(!cur) 
            {
                cur=new Node();
                cur->num=num;
                return;
            }
            
            if(num==cur->num)
            {
                cur->size++;
                return;
            }
            else if(num>cur->num)
            {
                insert(num,cur->rs);
                cur->size++;
            }
            else
            {
                insert(num,cur->ls);
                cur->size++;
            }
            
        }
        
        int queryLT(long long num, Node *cur)
        {
            if(cur==NULL) return 0;
            
            if(cur->num<num)
            {
                int temp=cur->size-((!cur->rs)?0:cur->rs->size);
                return queryLT(num,cur->rs)+temp;
            }
            else if(cur->num>num)
            {
                return queryLT(num,cur->ls);
            }
            
            else return cur->size-((!cur->rs)?0:cur->rs->size);
        }
        
    };

Log in to reply
 

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