Count of Range Sum LeetCode solution in Java

1
class Solution {
2
  public int countRangeSum(int[] nums, int lower, int upper) {
3
    int n = nums.length, ans = 0;
4
    long[] pre = new long[n + 1];
5
    for (int i = 0; i < n; i++) {
6
      pre[i + 1] = nums[i] + pre[i];
7
    }
8
    Arrays.sort(pre);
9
    int[] bit = new int[pre.length + 2];
10
    long sum = 0;
11
    for (int i = 0; i < n; i++) {
12
      update(bit, bs(sum, pre), 1);
13
      sum += nums[i];
14
      ans += sum(bit, bs(sum - lower, pre)) - sum(bit, bs(sum - upper - 1, pre));
15
    }
16
    return ans;
17
  }
18

19
  private int bs(long sum, long[] pre) { // return the index of first number bigger than sum
20
    int lo = 0, hi = pre.length;
21
    while (lo < hi) {
22
      int mid = (lo + hi) >> 1;
23
      if (pre[mid] > sum) {
24
        hi = mid;
25
      } else {
26
        lo = mid + 1;
27
      }
28
    }
29
    return lo;
30
  }
31

32
  private void update(int[] bit, int idx, int inc) {
33
    for (++idx; idx < bit.length; idx += idx & -idx) {
34
      bit[idx] += inc;
35
    }
36
  }
37

38
  private int sum(int[] bit, int idx) {
39
    int ans = 0;
40
    for (++idx; idx > 0; idx -= idx & -idx) {
41
      ans += bit[idx];
42
    }
43
    return ans;
44
  }
45
}