Java solution using two PriorityQueues


  • 9

    Almost the same idea of Find Median from Data Stream https://leetcode.com/problems/find-median-from-data-stream/

    1. Use two Heaps to store numbers. maxHeap for numbers smaller than current median, minHeap for numbers bigger than and equal to current median. A small trick I used is always make size of minHeap equal (when there are even numbers) or 1 element more (when there are odd numbers) than the size of maxHeap. Then it will become very easy to calculate current median.
    2. Keep adding number from the right side of the sliding window and remove number from left side of the sliding window. And keep adding current median to the result.
    public class Solution {
        PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(
            new Comparator<Integer>() {
                public int compare(Integer i1, Integer i2) {
                    return i2.compareTo(i1);
                }
            }
        );
    	
        public double[] medianSlidingWindow(int[] nums, int k) {
            int n = nums.length - k + 1;
    	if (n <= 0) return new double[0];
            double[] result = new double[n];
            
            for (int i = 0; i <= nums.length; i++) {
                if (i >= k) {
            	result[i - k] = getMedian();
            	remove(nums[i - k]);
                }
                if (i < nums.length) {
            	add(nums[i]);
                }
            }
            
            return result;
        }
        
        private void add(int num) {
    	if (num < getMedian()) {
    	    maxHeap.add(num);
    	}
    	else {
    	    minHeap.add(num);
    	}
    	if (maxHeap.size() > minHeap.size()) {
                minHeap.add(maxHeap.poll());
    	}
            if (minHeap.size() - maxHeap.size() > 1) {
                maxHeap.add(minHeap.poll());
            }
        }
    	
        private void remove(int num) {
    	if (num < getMedian()) {
    	    maxHeap.remove(num);
    	}
    	else {
    	    minHeap.remove(num);
    	}
    	if (maxHeap.size() > minHeap.size()) {
                minHeap.add(maxHeap.poll());
    	}
            if (minHeap.size() - maxHeap.size() > 1) {
                maxHeap.add(minHeap.poll());
            }
        }
    	
        private double getMedian() {
    	if (maxHeap.isEmpty() && minHeap.isEmpty()) return 0;
    	    
    	if (maxHeap.size() == minHeap.size()) {
    	    return ((double)maxHeap.peek() + (double)minHeap.peek()) / 2.0;
    	}
    	else {
                return (double)minHeap.peek();
    	}
        }
    }
    

  • 3

    Complexity is only O(nk), right?


  • 0

    @StefanPochmann I guess it's O(N lgK)? As each add() and remove() operation takes lgK time? And there are N elements to operate.
    getMedian() is O(1).


  • 0

    @shawngao The documentation says remove(Object) takes linear time.


  • 0

    @StefanPochmann You are right. Then the overall runtime complexity is O(nk).


  • 0
    C

    @shawngao It's O(nk), not O(nlogk).


  • 1
    D

    Time complexity is O(nlogk) if using Treemap instead of heap.


  • 0
    J

    @dtcxzch11 but you need to take care of the duplicates and it is really annoying


  • 1
    C

    Thank you @shawngao for this solution. I have made the code more concise by removing the duplicate parts.
    Thank you @StefanPochmann for your suggestions. I have included them.

    public class Solution{
        private PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();
        private PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(Collections.reverseOrder());
    	
        public double[] medianSlidingWindow(int[] nums, int k) {
            int n = nums.length - k + 1;
    	if (n <= 0) return new double[0];
    
            double[] result = new double[n];
            
            for (int i = 0; i <= nums.length; i++) {
                if (i >= k){
                	result[i - k] = getMedian();
                	remove(nums[i - k]);
                }
                if(i < nums.length){
        	        add(nums[i]);
                }
            }
            return result;
        }
        
        private void add(int num) {
        	(num < getMedian() ? maxHeap : minHeap).add(num);
        	resizeHeaps();
        }
    	
        private void remove(int num) {
        	(num < getMedian() ? maxHeap : minHeap).remove(num);
        	resizeHeaps();
        }
        
        private void resizeHeaps(){
            if (maxHeap.size() > minHeap.size()) {
                minHeap.add(maxHeap.poll());
        	}
            if (minHeap.size() - maxHeap.size() > 1) {
                maxHeap.add(minHeap.poll());
            }
        }
    	
        private double getMedian() {
        	if (maxHeap.isEmpty() && minHeap.isEmpty()) return 0;
        	else if (maxHeap.size() == minHeap.size()) {
        	    return ((double) maxHeap.peek() + (double) minHeap.peek()) / 2.0;
        	}
        	else {
                return (double) minHeap.peek();
        	}
        }
    }

  • 3

    @chaitanya0411 If you like concise, you could also do (num < getMedian() ? maxHeap : minHeap).add(num);.


  • 0

    Actually you can just implement minHeap and maxHeap by your self, and using a HashMap to access the position of need-to-removed node, in this way, you can achieve the time complexity in O(nlogk).


  • 0

    @jimhust1993 said in Java solution using two PriorityQueues:

    @dtcxzch11 but you need to take care of the duplicates and it is really annoying

    It's not too bad. See my solution using two Tree Sets.


  • -1
    T

    @shawngao Fails for input [1] , k = 1


  • 0

    @thepha3drus Really? It works for me.

    Custom Testcase
    [1]
    1
    
    Your answer
    [1.00000]
    
    Expected answer
    [1.00000]
    

  • 0
    X

    Well, this is O(nk), same as brute force method's time complexity. And brute force method even beats this on space complexity.


  • 0
    P

    My answer using TreeSet to implement heaps, similar to the answer of @shawngao.
    When writing the compareTo method in Node class, I discovered that if simply write code like this would not pass the last test case, because Integer.MIN_VALUE - Integer.MAX_VALUE returns 1 which is a postive value.
    It took some time for me to find the mistake, so I post this to prevent you guys from making the same mistake as me.

     public int compareTo(Node other) {
            if (this.val != other.val) {
                return this.val - other.val; 
            } else {
                return this.index - other.index;
            }
        }
    
    class Node implements Comparable<Node>{
        public int index;
        public int val;
        
        public Node(int index, int val) {
            this.index = index;
            this.val = val;
        }
        
        @Override
        public int compareTo(Node other) {
            if (this.val != other.val) {
                return new Integer(this.val).compareTo(new Integer(other.val)); 
            } else {
                return this.index - other.index;
            }
        }
    }
    
    
    public class Solution {
        
        TreeSet<Node> minHeap = new TreeSet<>();
        TreeSet<Node> maxHeap = new TreeSet<>();
        
        public double[] medianSlidingWindow(int[] nums, int k) {
            List<Double> list = new ArrayList<>();
            
            if (nums == null || nums.length == 0 || k > nums.length) {
                return new double[list.size()];
            }
            
            for (int i = 0; i < k-1; i++) {  //add k-1 numbers into two heaps
                add(new Node(i, nums[i]));
            }
            
            for (int i = k-1; i < nums.length; i++) {
                add(new Node(i, nums[i]));
                list.add(getMedian());
                remove(new Node(i-k+1, nums[i-k+1]));
            }
            
            double[] result = new double[list.size()];
            for (int i = 0; i < result.length; i++) {
                result[i] = list.get(i);
            }
            return result;
        }
        
        public void add(Node n) {
            maxHeap.add(n);
            minHeap.add(maxHeap.pollLast());
            
            if (maxHeap.size() < minHeap.size()) {
                maxHeap.add(minHeap.pollFirst());
            }
            
            if (maxHeap.size() - minHeap.size() > 1) {
                minHeap.add(maxHeap.pollLast());
            }
        }
        
        public void remove(Node n) {
            if (minHeap.contains(n)) {
                minHeap.remove(n);
            } else {
                maxHeap.remove(n);
            }
            
            if (maxHeap.size() < minHeap.size()) {
                maxHeap.add(minHeap.pollFirst());
            }
            
            if (maxHeap.size() - minHeap.size() > 1) {
                minHeap.add(maxHeap.pollLast());
            }
        }
        
        public double getMedian() {
            if (minHeap.size() == maxHeap.size()) {
                return minHeap.first().val/2.0 + maxHeap.last().val/2.0;
            } else {
                return maxHeap.last().val;
            }
        }
    }
    

  • 0
    G

    Same idea. Simple Version.

    public double[] medianSlidingWindow(int[] nums, int k) {
        double[] result = new double[nums.length - k + 1];
        PriorityQueue<Integer> left = new PriorityQueue<>(Collections.reverseOrder());
        PriorityQueue<Integer> right = new PriorityQueue<>();
        
        for(int i = 0; i < nums.length; i++) {
            if(left.size() <= right.size()) {
                right.add(nums[i]);
                left.add(right.remove());
            } else {
                left.add(nums[i]);
                right.add(left.remove());
            }
            
    
            if(left.size() + right.size() == k) {
                double median;
                if(left.size() == right.size()) {
                    median = (double) ((long)left.peek() + (long)right.peek()) / 2;
                } else {
                    median = (double) left.peek();
                }
                
                int start = i - k + 1;
                result[start] = median;
                if(!left.remove(nums[start])) {
                    right.remove(nums[start]);
                }
            }
        }
        return result;
    }

Log in to reply
 

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