Java Solution using two priority queues


  • 0
    M

    The idea is from the running median of a stream. We maintain two heaps, one is min heap and another one is max heap. Why we did this is because we want to split the elements in current window as two parts, the smaller part and the larger part. For example, if the window is [3,-1,5,2]. The two heaps will be: max heap A= [2, -1], min heap B= [3, 5], so the median of the window will be (A.peek() + B.peek())/2.0. We should always keep the size of two heap equal or the B heap will have just 1 more element than A heap. If any state violate this rule, we should repair the state.

    public class Solution {
        public double[] medianSlidingWindow(int[] nums, int k) {
            PriorityQueue<Integer> sq = new PriorityQueue<>();
            PriorityQueue<Integer> lq = new PriorityQueue<>((i1, i2) -> { // simply returning i2 - i1 could result in overflow issue.
                if (i1 > i2) {
                    return -1;
                } else if (i1 < i2) {
                    return 1;
                }
                return 0;
            });
            double[] res = new double[nums.length - k + 1];
            int c = 0;
            for (int i = 0; i < nums.length; i ++) {
                if (lq.size() == 0 || nums[i] >= lq.peek()) {
                    sq.offer(nums[i]);
                    if (sq.size() > lq.size() + 1) { // repair the state
                        lq.offer(sq.poll());
                    }
                } else {
                    lq.add(nums[i]);
                    if (lq.size() > sq.size()) { // repair the state
                        sq.offer(lq.poll());
                    }
                }
                if (i > k - 1) { // remove the elements out of window
                    int rm = nums[i - k]; 
                    if (lq.remove(rm)) {
                        if (lq.size() < sq.size() - 1) {
                            lq.offer(sq.poll());
                        }
                    } else {
                        sq.remove(rm);
                        if (lq.size() > sq.size()) {
                            sq.offer(lq.poll());
                        }
                    }
                }
                if (i >= k-1) {
                    res[c++] = k % 2 == 0 ? ((double)lq.peek() + (double)sq.peek())/2.0 : (double)sq.peek();
                }
            }
            return res;
        }
    }
    

    Please help to refine the algorithm and make it more concise. And look forward to seeing any better code!


  • 0
    S

    Unfortunately, PriorityQueue.remove(Object o) is not an O(logn) operation, it is O(n).


  • 0
    M

    @si-yao Thank you for pointing out! I agree with you. So if we want to achieve O(n*log(k) with this solution, we could use TreeMap to replace PriorityQueue (TreeSet is not okay because duplicate elements are allowed, we need to use TreeMap to record the count of elements). Then other stuff are pretty the same logic with priority queue although the code might be a little bit verbose when we use tree map.


Log in to reply
 

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