Short and Clear O(nlogk) Java Solutions


  • 3

    Use 2 priorityQueues:
    3 steps: remove out-of-bound element, add new element and keep small.peek()<=big.peek(), and get median.
    Time: O(nk). The drawback is remove function of priorityQueue takes O(k) time.

    public class Solution {
        public double[] medianSlidingWindow(int[] nums, int k) {
            double[] res = new double[nums.length - k + 1];
            int idx = 0;
            boolean useBoth = k % 2 == 0;
            PriorityQueue<Integer> small = new PriorityQueue<>((a, b)->{return (int)((double)b-a);});
            PriorityQueue<Integer> big = new PriorityQueue<>();
            for(int i = 0; i<nums.length; i++){
                if(small.size() + big.size()==k){
                    Integer toRemove = new Integer(nums[i-k]);
                    if(toRemove <= small.peek()) small.remove(toRemove);
                    else big.remove(toRemove);
                }
                //always keep small.size() == big.size() or small.size() == big.size()+1
                if(small.size()<=big.size()) small.add(nums[i]);
                else big.add(nums[i]);
                if(big.size()>0){
                    while(small.peek()>big.peek()){
                        big.add(small.poll());
                        small.add(big.poll());
                    }
                }
                if(small.size() + big.size()==k){
                    if(useBoth) res[idx++] = ((double)small.peek() + big.peek())/2.0;
                    else res[idx++] = (double)small.peek();
                }
            }
            return res;
        }
    }
    

    To overcome priorityQueue's remove O(k) problem and make our solution O(nlogk), we can replace head/priorityQueue to BST, which is treemap, as below. This would be complicated to write, because we need to deal with duplicated elements and update counts, but the logic is entirely the same as the above solution.
    Time: O(nlogk).

    public class Solution {
        public double[] medianSlidingWindow(int[] nums, int k) {
            double[] res = new double[nums.length - k + 1];
            int idx = 0;
            boolean useBoth = k % 2 == 0;
            TreeMap<Integer, Integer> small = new TreeMap<>((a, b)->{return (int)((double)b-a);});
            int smallSize = 0; 
            TreeMap<Integer, Integer> big = new TreeMap<>();
            int bigSize = 0;
            for(int i = 0; i<nums.length; i++){
                if(smallSize + bigSize == k){
                    if(nums[i-k] <= small.firstKey()){
                        remove(small, nums[i-k]);
                        smallSize--;
                    }else{
                        remove(big, nums[i-k]);
                        bigSize--;
                    }
                }
    
                if(smallSize<=bigSize){
                    add(small, nums[i]);
                    smallSize++;
                }else{
                    add(big, nums[i]);
                    bigSize++;
                }
                if(bigSize>0){
                    while(small.firstKey()>big.firstKey()){
                        add(big, remove(small, small.firstKey()));
                        add(small, remove(big, big.firstKey()));
                    }
                }
                
                if(smallSize + bigSize==k){
                    if(useBoth) res[idx++] = ((double)small.firstKey() + big.firstKey())/2.0;
                    else res[idx++] = (double)small.firstKey();
                }
            }
            return res;
        }
        
        private int remove(TreeMap<Integer, Integer> map, int i){
            map.put(i, map.get(i)-1);
            if(map.get(i)==0) map.remove(i);
            return i; 
        }
        
        private void add(TreeMap<Integer, Integer> map, int i){
            if(!map.containsKey(i)) map.put(i, 1);
            else map.put(i, map.get(i)+1);
        }
    }
    

  • 0
    Y

    This post is super clear and easy to understand!!
    Just a little improvement for the helper function which is more concise and might be faster.
    The searching process in remove part is only one time.

    private int remove(TreeMap<Integer, Integer> map, int i) {
    	Integer count = map.get(i);
    	if(count == 1) map.remove(i);
    	else map.put(i, count-1);
    	return i;
    }
    private void add(TreeMap<Integer,Integer> map, int i) {
    	map.put(i, map.getOrDefault(i, 0)+1);
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.