Sliding Window Median

Heap, Design, Two Pointer, Multi-set

Hard

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

Examples:

[2,3,4], the median is3

[2,3], the median is(2 + 3) / 2 = 2.5

Given an arraynums, there is a sliding window of sizekwhich is moving from the very left of the array to the very right. You can only see theknumbers in the window. Each time the sliding window moves right by one position. Your job is to output the median array for each window in the original array.

For example, Givennums=[1,3,-1,-3,5,3,6,7], andk= 3.

Window position                Median
---------------               -----
[1  3  -1] -3  5  3  6  7       1
 1 [3  -1  -3] 5  3  6  7       -1
 1  3 [-1  -3  5] 3  6  7       -1
 1  3  -1 [-3  5  3] 6  7       3
 1  3  -1  -3 [5  3  6] 7       5
 1  3  -1  -3  5 [3  6  7]      6

Therefore, return the median sliding window as[1,-1,-1,3,5,6].

Note: You may assumekis always valid, ie:kis always smaller than input array's size for non-empty array.

Solution & Analysis

Two Heaps - Min Heap + Max Heap - O(n*k) time, O(k) space

26 ms, faster than 96.12%

和 Find Median from Data Stream 类似思路，用两个Heap，一个Min-Heap，一个Max-Heap，这样Min-Heap, Max-Heap相当于各自维护了k/2的数字，这里允许 maxHeap.size() == minHeap.size() + 1，也就是说在k为奇数情况下，允许maxHeap中多一个元素，这样median就是maxHeap.peek()，否则k为偶数时，是 (minHeap.peek() + maxHeap.peek()) / 2。

这里注意溢出问题，将peek()的值强制转换为double可以防止在相加/2时溢出。

@dico: PriorityQueue maxHeap = new PriorityQueue<>((a,b)->(b-a)) won't pass test cases that involve comparing Integer MAX and MIN, change it to: PriorityQueue maxHeap = new PriorityQueue<>((a,b)->(b.compareTo(a)));

Implementation - Two priority queues:

  1. A max-heap to store the smaller half of the numbers

  2. A min-heap to store the larger half of the numbers

  3. If k = 2*n + 1 (∀n∈Z), then max-heap is allowed to hold n+1 elements, while min-heap can hold n elements.

  4. If k = 2*n (∀n∈Z), then both heaps are balanced and hold nn elements each.

Complexity

  • Time:

    • add() - O(logk)

    • remove() - O(logk + k) ~ O(k). Due to remove(Object o) is O(k)

    • total - O(n * k)

  • Space: O(k)

class Solution {
    private PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    private PriorityQueue<Integer> maxHeap = new PriorityQueue<>(new Comparator<Integer>() {
        @Override
        public int compare(Integer a, Integer b) {
            return Integer.compare(b, a);
        }
    });
    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()) {
            minHeap.add(num);
        } else {
            maxHeap.add(num);
        }
        rebalance();
    }

    private void remove(int num) {
        if (num > getMedian()) {
            minHeap.remove(num);
        } else {
            maxHeap.remove(num);
        }
        rebalance();
    }

    private double getMedian() {
        if (minHeap.isEmpty() && maxHeap.isEmpty()) {
            return 0;
        } 
        if (minHeap.size() == maxHeap.size()) {
            return ((double) minHeap.peek() + (double) maxHeap.peek()) / 2.0;
        }
        return (double) maxHeap.peek();
    }

    private void rebalance() {
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.add(minHeap.poll());
        } else if (maxHeap.size() - minHeap.size() > 1) {
            minHeap.add(maxHeap.poll());
        }
    }
}

Similar Implementation - Custom MedianQueue class using Two Heaps

Inspired by @yuxiangmusic: https://leetcode.com/problems/sliding-window-median/discuss/96339/Java-clean-and-easily-readable-solution-with-a-helper-class

class Solution {
    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];
        MedianQueue window = new MedianQueue();

        for (int i = 0; i <= nums.length; i++) {
            if (i >= k) {
                result[i - k] = window.getMedian();
                window.remove(nums[i - k]);
            }
            if (i < nums.length) {
                window.add(nums[i]);
            }
        }
        return result;
    }
    class MedianQueue {

        private PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        private PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

        public void add(int num) {
            if (num > getMedian()) {
                minHeap.add(num);
            } else {
                maxHeap.add(num);
            }
            rebalance();
        }

        public void remove(int num) {
            if (num > getMedian()) {
                minHeap.remove(num);
            } else {
                maxHeap.remove(num);
            }
            rebalance();
        }

        public double getMedian() {
            if (minHeap.isEmpty() && maxHeap.isEmpty()) {
                return 0;
            } 
            if (minHeap.size() == maxHeap.size()) {
                return ((double) minHeap.peek() + (double) maxHeap.peek()) / 2.0;
            }
            return (double) maxHeap.peek();
        }

        public void rebalance() {
            if (minHeap.size() > maxHeap.size()) {
                maxHeap.add(minHeap.poll());
            } else if (maxHeap.size() - minHeap.size() > 1) {
                minHeap.add(maxHeap.poll());
            }
        }
    }
}

更优化的方案是使用Hash-Heap，也就是优化remove(Object o)这个O(n)的操作. 但是Java没有内置Hash-Heap，可以用TreeMap或者TreeSet来实现

TreeMap

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);
    }
}

TreeSet

Time: O(nlogk)

23 ms

public class Solution {
    class MedianSet {
        private TreeSet<Integer> left;
        private TreeSet<Integer> right;

        public MedianSet(int[] nums) {
            Comparator<Integer> cmp = new Comparator<Integer>(){
                public int compare(Integer a, Integer b) {
                    return nums[a] == nums[b] ? a - b : nums[a] < nums[b] ? -1 : 1;
                }
            };
            left = new TreeSet<>(cmp);
            right = new TreeSet<>(cmp);
        }

        public void add(int i) {
            if (left.size() <= right.size()) {
                right.add(i);
                int m = right.first();
                right.remove(m);
                left.add(m);
            } else {
                left.add(i);
                int m = left.last();
                left.remove(m);
                right.add(m);
            }

        }

        public void remove(int i) {
            if(!left.remove(i)) {
                right.remove(i);
            }
        }

        public double getMedian(int[] nums) {
            double median;
            if (left.size() == right.size()) {
                median = ((double) nums[left.last()] + nums[right.first()]) / 2;
            } else if (left.size() < right.size()) {
                median = nums[right.first()];
            } else {
                median = nums[left.last()];
            }

            return median;
        }

    }

    public double[] medianSlidingWindow(int[] nums, int k) {
        int n = nums.length - k + 1;
        double[] result = new double[n];
        MedianSet window = new MedianSet(nums);
        for(int i = 0; i <= nums.length; i++) {
            if (i >= k) {
                result[i - k] = window.getMedian(nums);
                window.remove(i - k);
            }
            if (i < nums.length) {
                window.add(i);
            }
        }
        return result;
    }

}

Reference

https://leetcode.com/problems/sliding-window-median/solution/

PreviousFind Median from Data StreamNextDesign Hit Counter

Last updated