0327.cpp (1687B)
1 class SegmentTree { 2 const int n; 3 vector<long long> t = vector<long long>(2 * n); 4 5 public: 6 SegmentTree(int n) : n(n) {} 7 8 int sum(int l, int r) const { 9 int res = 0; 10 11 for (l += n, r += n; l < r; l /= 2, r /= 2) { 12 if (l & 1) res += t[l++]; 13 if (r & 1) res += t[--r]; 14 } 15 16 return res; 17 } 18 19 void update(int p) { 20 for (t[p += n]++; p > 1; p /= 2) 21 t[p / 2] = t[p] + t[p ^ 1]; 22 } 23 }; 24 25 class Solution { 26 public: 27 int countRangeSum(const vector<int> &nums, int lower, int upper) const { 28 const int n = size(nums); 29 vector<long long> prefix(n + 1); 30 31 for (int i = 0; i < n; i++) { 32 prefix[i + 1] = 1ll * prefix[i] + nums[i]; 33 } 34 35 vector<long long> possible; 36 for (const long long n : prefix) { 37 possible.push_back(n); 38 possible.push_back(n - lower); 39 possible.push_back(n - upper); 40 } 41 42 const int m = size(possible); 43 static int idxs[300004]; 44 iota(idxs, idxs + m, 0); 45 sort(idxs, idxs + m, [&](int a, int b) { return possible[a] < possible[b]; }); 46 47 static int rev[300004]; 48 for (int i = m - 2, cnt = rev[idxs[m - 1]] = m - 1; i >= 0; i--) { 49 if (possible[idxs[i]] != possible[idxs[i + 1]]) cnt--; 50 rev[idxs[i]] = cnt; 51 } 52 53 int res = 0; 54 SegmentTree seg(size(possible)); 55 for (int i = 0; i <= n; i++) { 56 const int low = rev[i * 3 + 2]; 57 const int high = rev[i * 3 + 1]; 58 res += seg.sum(low, high + 1); 59 seg.update(rev[i * 3]); 60 } 61 62 return res; 63 } 64 };