BiruLyu
7/5/2017 - 10:42 PM

560. Subarray Sum Equals K(1st).java

public class Solution {
    public int subarraySum(int[] nums, int k) {
        if (nums == null || nums.length < 1) return 0;
        Map<Integer, Integer> map = new HashMap<>();
        int prefixSum = 0;
        map.put(0, 1);
        int res = 0;
        for (int num : nums) {
            prefixSum += num;
            int diff = prefixSum - k;
            res += map.getOrDefault(diff, 0);
            map.put(prefixSum, map.getOrDefault(prefixSum, 0) + 1);
        }
        return res;
    }
}
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;
    }
}