Count of Range Sum LeetCode solution in C++

1
class Solution {
2
 public:
3
  using ll = long long;
4

5
  int countRangeSum(vector<int> &nums, int lower, int upper) {
6
    // Build prefix sums
7
    vector<ll> prefixSums(nums.size() + 1, 0);
8
    for (int i = 0; i < nums.size(); i++) {
9
      prefixSums[i + 1] = prefixSums[i] + nums[i];
10
    }
11

12
    // Run merge sort and count range sum along the way
13
    tempNums.assign(prefixSums.size(), 0);
14
    splitAndMerge(prefixSums, 0, prefixSums.size(), lower, upper);
15

16
    return count;
17
  }
18

19
  void splitAndMerge(vector<ll> &nums, const int left, const int right, const int lower,
20
                     const int upper) {
21
    if (right - left <= 1) return;
22
    const int mid = left + (right - left) / 2;
23
    splitAndMerge(nums, left, mid, lower, upper);
24
    splitAndMerge(nums, mid, right, lower, upper);
25

26
    countRangeSums(nums, left, mid, right, lower, upper);
27

28
    merge(nums, left, mid, right);
29
  }
30

31
  void countRangeSums(const vector<ll> &prefixSums, const int left, const int mid, const int right,
32
                      const int lower, const int upper) {
33
    // S(i,j) == prefixSums[j+1] - prefixSums[i] (i <= j)
34
    // S(i,j) == prefixSums[k] - prefixSums[i] (let k=j+1, i < k)
35
    //
36
    // lower <= S(i,j) <= upper
37
    // => lower <= prefixSums[k] - prefixSums[i] <= upper
38
    // => lower + prefixSums[i] <= prefixSums[k] <= upper + prefixSums[i]
39
    for (int i = left; i < mid; i++) {
40
      const ll newLower = lower + prefixSums[i];
41
      const ll newUpper = upper + prefixSums[i];
42
      const auto findStart = prefixSums.begin() + mid;
43
      const auto findEnd = prefixSums.begin() + right;
44
      const auto itFoundLower = std::lower_bound(findStart, findEnd, newLower);
45
      const auto itFoundUpper = std::upper_bound(findStart, findEnd, newUpper);
46
      count += (itFoundUpper - itFoundLower);
47
    }
48
  }
49

50
  void merge(vector<ll> &nums, const int left, const int mid, const int right) {
51
    std::copy(nums.begin() + left, nums.begin() + right, tempNums.begin() + left);
52
    int i = left, j = mid;
53
    for (int k = left; k < right; k++) {
54
      if (i < mid && j < right) {
55
        if (tempNums[i] < tempNums[j]) {
56
          nums[k] = tempNums[i++];
57
        } else {
58
          nums[k] = tempNums[j++];
59
        }
60
      } else if (i >= mid) {
61
        nums[k] = tempNums[j++];
62
      } else { // j >= right
63
        nums[k] = tempNums[i++];
64
      }
65
    }
66
  }
67

68
  vector<ll> tempNums;
69
  int count = 0;
70
};