327 Count of Range Sum

327. Count of Range Sum

1. Question

Given an integer arraynums, return the number of range sums that lie in[lower, upper]inclusive. Range sumS(i, j)is defined as the sum of the elements innumsbetween indicesiandj(i≤j), inclusive.

Note: A naive algorithm of O(n^2) is trivial. You MUST do better than that.

Example: Givennums=[-2, 5, -1],lower=-2,upper=2, Return3. The three ranges are :[0, 0],[2, 2],[0, 2]and their respective sums are:-2, -1, 2.

2. Implementation

(1) TreeMap

我们定义sum[i]为nums[0...i]之间的subarray sum, 对于区间[j,i]的subarray sum，其值等于sum[i] - sum[j], 如果 lower <= sum[i] - sum[j] <= upper, 通过推导可以得出 sum[j] - upper <= sum[j] <= sum[i] - lower. 我们遍历数组的时候，利用sum变量加上当前的数，这样我们很容易得到sum[i], 要得到sum[j], 我们利用TreeMap的区间查询的特性，查询treemap，key在 sum[i] - upper[i] 和 sum[i] - lower之间的值的个数，最后得到结果

public class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {        
        TreeMap<Long, Long> map = new TreeMap();
        map.put((long)0, (long)1);
        long sum = 0;
        int count = 0;
        for (int i = 0; i < nums.length; i++) {
            sum += nums[i];
            long fromKey = sum - upper;
            long toKey = sum - lower;
            Map<Long, Long> subMap = map.subMap(fromKey, true, toKey, true);

            for (long value : subMap.values()) {
                count += value;
            }

            if (map.containsKey(sum)) {
                map.put(sum, map.get(sum) + 1);
            }
            else {
                map.put(sum, (long)1);
            }
        }
        return count;
    }
}

(2) Merge Sort

思路: 这里的思路和Leetcode 315很类似，我们利用merge sort，在merge的过程中count 符合条件的区间和个数. 具体在我们的merge()里，我们知道sum[start, mid] 和 sum[mid + 1, end]这两个区间是已经排好序的，对于每个在左区间的数sum[left], 我们分别要找出一个index high满足sum[high] - sum[left] > upper, 以及第一个index low满足sum[low] - sum[left] >= lower, 这样符合条件的区间个数就是 high - low. 在计算符合条件的区间个数同时，我们也顺便完成merge的过程

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        if (nums == null || nums.length == 0) {
            return 0;
        }

        int n = nums.length;
        long[] sum = new long[n + 1];

        for (int i = 1; i <= n; i++) {
            sum[i] = sum[i - 1] + nums[i - 1];
        }

        int[] count = new int[1];
        mergeSort(sum, 0, sum.length - 1, lower, upper, count);
        return count[0];
    }

    public void mergeSort(long[] sum, int start, int end, int lower, int upper, int[] count) {
        if (start >= end) {
            return;
        }

        int mid = start + (end - start) / 2;
        mergeSort(sum, start, mid, lower, upper, count);
        mergeSort(sum, mid + 1, end, lower, upper, count);
        merge(sum, start, mid, end, lower, upper, count);
    }

    public void merge(long[] sum, int start, int mid, int end, int lower, int upper, int[] count) {
        long[] temp = new long[end - start + 1];

        int low = mid + 1;
        int high = mid + 1;
        int right = mid + 1;
        int tIndex = 0;

        for (int left = start; left <= mid; left++) {
            while (low <= end && sum[low] - sum[left] < lower) ++low;
            while (high <= end && sum[high] - sum[left] <= upper) ++high;

            count[0] += high - low;

            while (right <= end && sum[right] <= sum[left]) temp[tIndex++] = sum[right++];
            temp[tIndex++] = sum[left];
        }

        while (right <= end) {
            temp[tIndex++] = sum[right++];
        }

        for (int k = start; k <= end; k++) {
            sum[k] = temp[k - start];
        }
    }
}

3. Time & Space Complexity

TreeMap: 时间复杂度O(nlogn), 空间复杂度O(n)

Merge Sort: 时间复杂度O(nlogn), 空间复杂度O(n)