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].

思路

因为核心是subarray sum equals k问题，所以使用prefixSum求解。但是找的是cloeset，所以使用TreeMap instead of HashMap

复杂度

time O(nlogn), space O(n)

代码

  public class Solution {
    /**
     * @param nums: A list of integers
     * @return: A list of integers includes the index of the first number 
     *          and the index of the last number
     */
    public int[] subarraySumClosest(int[] nums) {
        // write your code here
        // check edge case
        if (nums == null || nums.length == 0) {
            return new int[0];
        }
        // preprocess
        TreeMap<Integer, Integer> map = new TreeMap<Integer, Integer>();
        map.put(0, -1);
        int diff = Integer.MAX_VALUE;
        int[] result = new int[]{0,0};
        int prefix = 0;
        // main
        for (int i = 0; i < nums.length; i++) {
            prefix += nums[i];

            //floor
            if (map.floorKey(prefix) != null) {
                int floor = map.floorKey(prefix);
                int floorDiff = Math.abs(prefix - floor);

                if (floorDiff < diff) {
                    diff = floorDiff;
                    result[0] = map.get(floor) + 1;
                    result[1] = i;
                }
            }

            //ceiling
            if (map.ceilingKey(prefix) != null) {
                int ceiling = map.ceilingKey(prefix);
                int ceilingDiff = Math.abs(prefix - ceiling);
                if (ceilingDiff < diff) {
                    diff = ceilingDiff;
                    result[0] = map.get(ceiling) + 1;
                    result[1] = i;
                }
            }

            map.put(prefix, i);
        }
        return result;
    }
}