Lintcode 139. Subarray Sum Closest

Given an integer array, find a subarray with sum closest to zero. Return the indexes of the first number and last number.

Example

Given [-3, 1, 1, -3, 5], return [0, 2], [1, 3], [1, 1], [2, 2] or [0, 4].

Challenge

O(nlogn) time

Solution:
The naive solution is for each presum[j], check out all its previous presum[i -1] and calculate the closest one. The overall time complexity is O(n^2).

A better solution is to use a TreeMap, where elements are sorted in the tree. Then for each presum[j], we only need to find out the closest element in the tree, which takes time of O(logn). So the overall time complexity is O(nlogn).

Code (java):

public class Solution {
    /**
     * @param A: an integer array
     * @param target: An integer
     * @param k: An integer
     * @return: an integer array
     */
    public int[] kClosestNumbers(int[] A, int target, int k) {
        // write your code here
        if (A == null || A.length == 0) {
            return new int[0];
        }

        int[] ans = new int[k];

        // step 1: find the position of the largest number smaller than the target
        //
        int start = findCloestNumber(A, target);
        int end = start + 1;

        // step 2: find the k-cloest numbers
        //
        for (int i = 0; i < k; i++) {
            if (start < 0) {
                ans[i] = A[end];
                end++;
            } else if (end >= A.length) {
                ans[i] = A[start];
                start--;
            } else if (Math.abs(A[start] - target) <= Math.abs(A[end] - target)) {
                ans[i] = A[start];
                start--;
            } else {
                ans[i] = A[end];
                end++;
            }
        }

        return ans;
    }
    // find the first position of the closet number which is smaller than the target
    //
    private int findCloestNumber(int[] A, int target) {
        int start = 0;
        int end = A.length - 1;

        while (start + 1 < end) {
            int mid = start + (end - start) / 2;
            if (A[mid] == target) {
                end = mid;
            } else if (A[mid] > target) {
                end = mid;
            } else {
                start = mid;
            }
        }

        if (A[end] < target) {
            return end;
        }

        if (A[start] < target) {
            return start;
        }

        return -1;
    }
}