public class Solution {
public int subarraySum(int[] nums, int k) {
int[] sums = new int[nums.length + 1];
for (int i = 1; i < sums.length; i++) {
sums[i] = sums[i - 1] + nums[i - 1];
}
return sub(sums, new int[sums.length], 0, sums.length - 1, k);
}
private int sub(int[] sums, int[] aux, int start, int end, int k) {
if (start >= end) {
return 0;
}
int mid = start + (end - start) / 2;
int res = sub(sums, aux, start, mid, k) + sub(sums, aux, mid + 1, end, k);
res += merge(sums, aux, start, mid, end, k);
return res;
}
private int merge(int[] sums, int[] aux, int start, int mid, int end, int k) {
// System.out.println("start: " + start + " mid: " + mid + " end: " + end);
// System.out.println("Before merge: " + Arrays.toString(sums));
for (int i = start; i <= end; i++) {
aux[i] = sums[i];
}
int count = 0;
int i = start, j = mid + 1, t = start;
int p = mid + 1;
while (i <= mid) {
while (p <= end && aux[p] - k < aux[i]) p++;
for (int q = 0; q + p <= end; q++) {
if (aux[p + q] - k > aux[i]) break;
else count++;
}
while (j <= end && aux[i] > aux[j]) sums[t++] = aux[j++];
sums[t++] = aux[i++];
}
while (j <= end) sums[t++] = aux[j++];
// System.out.println("After merge: " + Arrays.toString(sums));
return count;
}
}