Lintcode 404. Subarray Sum II

题目

Given an integer array, find a subarray where the sum of numbers is in a given interval. Your code should return the number of possible answers. (The element in the array should be positive)

Example Given [1,2,3,4] and interval = [1,3], return 4. The possible answers are:

[0, 0] [0, 1] [1, 1] [2, 2]

思路

这道题如果用平常的prefixSum + HashMap/TreeMap去解会发现复杂度非常高，和暴力解并没有区别。这里我们要做的其实是range query。如果用TreeMap(BST)去做这个事，复杂度比较高，因为Map不能存key相同的元素，所以Map只能存成sum->frequency。我们需要遍历Map来得到区间内每一个key的count，worst case会是O(n)。

但是如果我们是在sorted array里面做range query,我们可以存的下每一个duplicate，这样我们可以通过index直接得出range内元素的个数。

这样这道题就转化为number of elements within a range in a sorted array with duplicates。我们要做的就是找到first element >= start的index，以及first element >= end + 1的index，做减法就可以了。注意为乐得到sorted array数组，我们需要先去compute一个prefixSum数组并对其排序。这里我们只需要count不需要index，所以排序并不会使我们丢失信息。

复杂度

time O(nlogn) because of sorting, space O(1)

代码

  public class Solution {
    /**
     * @param A an integer array
     * @param start an integer
     * @param end an integer
     * @return the number of possible answer
     */
    public int subarraySumII(int[] A, int start, int end) {
        // Write your code here
        // check edge
        if (A == null || A.length == 0 || start > end) {
            return 0;
        }
        // preprocess
        for (int i = 1; i < A.length; i++) {
            A[i] += A[i-1];
        }
        Arrays.sort(A);

        // main
        int count = 0;
        for (int i = 0; i < A.length; i++) {
            int cur = A[i];
            if (cur >= start && cur <= end) {
                count++;
            }

            int left = cur - end;
            int right = cur - start;
            count += (binarySearch(A, right + 1) - binarySearch(A, left));
        }
        return count;
    }
    //find first occurance of the target
    public int binarySearch(int[] A, int target) {
        if (A[A.length - 1] < target) {
            return A.length;
        }
        int start = 0;
        int end = A.length - 1;
        int mid;
        while (start + 1 < end) {
            mid = (start + end)/2;
            if (A[mid] == target) {
                end = mid;
            } else if (A[mid] > target) {
                end = mid;
            } else {
                start = mid;
            }
        }

        //budao
        if (A[start] >= target) {
            return start;
        }
        return end;
    }
}