leetcode

Solution to some Leetcode problems written in C++
git clone git://git.dimitrijedobrota.com/leetcode.git
Log | Files | Refs | README | LICENSE

0327.cpp (1687B)


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