leetcode

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 };