3 solutions: divide and conquer / order statistics tree / Fenwick tree


  • 2
    M
    // Count of Range Sum
    
    // Fenwick tree
    class Solution {
      void add(vector<int> &fenwick, int n, int x) {
        for (; x < n; x |= x+1)
          fenwick[x]++;
      }
      int getSum(vector<int> &fenwick, int x) {
        int s = 0;
        for (; x; x &= x-1)
          s += fenwick[x-1];
        return s;
      }
    public:
      int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size();
        long s = 0, cnt = 0;
        vector<int> fenwick(n+1, 0);
        vector<long> scale;
        scale.push_back(0);
        for (int x: nums)
          scale.push_back(s += x);
        sort(scale.begin(), scale.end());
        s = 0;
        add(fenwick, n+1, lower_bound(scale.begin(), scale.end(), 0) - scale.begin());
        for (int x: nums) {
          s += x;
          cnt += getSum(fenwick, upper_bound(scale.begin(), scale.end(), s-lower) - scale.begin()) -
            getSum(fenwick, lower_bound(scale.begin(), scale.end(), s-upper) - scale.begin());
          add(fenwick, n+1, lower_bound(scale.begin(), scale.end(), s) - scale.begin());
        }
        return cnt;
      }
    };
    
    /// divide and conquer
    
    class Solution {
      vector<long> a, b;
      int lower, upper;
      int f(int l, int h) {
        if (h-l <= 1)
          return 0;
        int m = l+h >> 1, cnt = f(l, m) + f(m, h), i = l, j = l;
        for (int k = m; k < h; k++) {
          while (i < m && a[i] < a[k]-upper)
            i++;
          while (j < m && a[j] <= a[k]-lower)
            j++;
          cnt += j-i;
        }
        merge(a.begin()+l, a.begin()+m, a.begin()+m, a.begin()+h, b.begin()+l);
        copy_n(b.begin()+l, h-l, a.begin()+l);
        return cnt;
      }
    public:
      int countRangeSum(vector<int>& nums, int lower, int upper) {
        this->lower = lower;
        this->upper = upper;
        int n = nums.size();
        long s = 0;
        a.clear();
        a.push_back(0);
        for (int x: nums)
          a.push_back(s += x);
        b.resize(n+1);
        return f(0, n+1);
      }
    };
    
    /// order statistics tree
    
    #include <ext/pb_ds/assoc_container.hpp>
    #include <ext/pb_ds/tree_policy.hpp>
    using namespace __gnu_pbds;
    
    class Solution {
    public:
      int countRangeSum(vector<int>& nums, int lower, int upper) {
        typedef pair<long, int> pli;
        tree<pli, null_type, less<pli>, rb_tree_tag, tree_order_statistics_node_update> t;
        t.insert({0, 0});
        int id = 1;
        long s = 0, cnt = 0;
        for (int x: nums) {
          s += x;
          cnt += t.order_of_key({s-lower, id}) - t.order_of_key({s-upper, 0});
          t.insert({s, id++});
        }
        return cnt;
      }
    };
    

    https://github.com/MaskRay/LeetCode/blob/master/count-of-range-sum.cc


  • 0
    This post is deleted!

  • 0
    S

    Thanks. Good to know the pb_ds classes :)


  • 0
    H

    __gnu_pbds is banned now :(


Log in to reply
 

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