leetcode@ [327] Count of Range Sum (Binary Search)

Posted 2020-07-31 流白

https://leetcode.com/problems/count-of-range-sum/

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.
Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

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

Example:
Given nums = [-2, 5, -1], lower = -2, upper = 2,
Return 3.
The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

 

class pair {
    public int idx;
    public long val;
    public pair(int idx, long val) {
        super();
        this.idx = idx;
        this.val = val;
    }
    
}

class pairComparator implements Comparator {
    
    public int compare(Object o1, Object o2) {
        pair p1 = (pair) o1;
        pair p2 = (pair) o2;
        if(p1.val < p2.val) {
            return -1;
        } else {
            return 1;
        }
    }
    
}

public class Solution {
    
    public static ArrayList<pair> toSortedList(long[] sum) {
        
        ArrayList<pair> ls = new ArrayList<pair> ();
        for(int i=0; i<sum.length; ++i) {
            pair p = new pair(i, sum[i]);
            ls.add(p);
        }
        Collections.sort(ls, new pairComparator());
        return ls;
    }
    
    public static int binarySearch(ArrayList<pair> ls, int l, int r, long lb, long ub, int index) {
        
        if(l > r) {
            return 0;
        }
        if(l == r) {
            if(ls.get(l).val >= lb && ls.get(l).val <= ub && ls.get(l).idx >= index) {
                //System.out.println("candidate index range:  [" + index + ", " + ls.get(l).idx + "]");
                return 1;
            }
            return 0;
        }
        
        int rs = 0;
        
        int mid = (l + r) / 2;
        if(ls.get(mid).val < lb) {
            rs = binarySearch(ls, mid+1, r, lb, ub, index);
        } else if(ls.get(mid).val > ub) {
            rs = binarySearch(ls, l, mid-1, lb, ub, index);
        } else {
            rs = binarySearch(ls, l, mid-1, lb, ub, index) + binarySearch(ls, mid+1, r, lb, ub, index);
            if(ls.get(mid).idx >= index) {
                //System.out.println("candidate index range:  [" + index + ", " + ls.get(l).idx + "]");
                rs++;
            }
        }
        
        return rs;
    }
    
    public int countRangeSum(int[] nums, int lower, int upper) {
        
        int n = nums.length;
        if(n == 0) {
            return 0;
        }
        
        long[] sum = new long[n];
        sum[0] = nums[0];
        for(int i=1; i<n; ++i) {
            sum[i] = sum[i-1] + nums[i];
        }
        
        int rs = 0;
        ArrayList<pair> ls = toSortedList(sum);
        for(int i=0; i<n; ++i) {
            long new_lower = (long)lower + sum[i] - (long)nums[i];
            long new_upper = (long)upper + sum[i] - (long)nums[i];
            int count = binarySearch(ls, 0, ls.size()-1, new_lower, new_upper, i);
            rs += count;
        }
        
        return rs;
    }
}

View Code