// Count of Range Sum // Fenwick tree class Solution { void add(vector &fenwick, int n, int x) { for (; x < n; x |= x+1) fenwick[x]++; } int getSum(vector &fenwick, int x) { int s = 0; for (; x; x &= x-1) s += fenwick[x-1]; return s; } public: int countRangeSum(vector& nums, int lower, int upper) { int n = nums.size(); long s = 0, cnt = 0; vector fenwick(n+1, 0); vector 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 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& 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 #include using namespace __gnu_pbds; class Solution { public: int countRangeSum(vector& nums, int lower, int upper) { typedef pair pli; tree, 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; } };