310/310

Author: MaskRay

README.md

@@ -2,7 +2,7 @@
 
 [LeetCode solutions](http://maskray.me/blog/2014-06-29-leetcode-solutions) gives some thoughts on selected problems.
 
-Solved 309/309 problems.
+Solved 310/310 problems.
 
 ## Database
 
@@ -12,6 +12,7 @@ See [database.md](database.md)
 
 | # | Title | Solution |
 |---| ----- | -------- |
+|327|[Count of Range Sum](https://leetcode.com/problems/count-of-range-sum/)|[count-of-range-sum.cc](count-of-range-sum.cc)|
 |326|[Power of Three](https://leetcode.com/problems/power-of-three/)|[power-of-three.cc](power-of-three.cc)|
 |325|[Maximum Size Subarray Sum Equals k](https://leetcode.com/problems/maximum-size-subarray-sum-equals-k/)|[maximum-size-subarray-sum-equals-k.cc](maximum-size-subarray-sum-equals-k.cc)|
 |324|[Wiggle Sort II](https://leetcode.com/problems/wiggle-sort-ii/)|[wiggle-sort-ii.cc](wiggle-sort-ii.cc)|

count-of-range-sum.cc

@@ -0,0 +1,93 @@
+// 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;
+  }
+};