[LeetCode 480] Sliding Window Median

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 is 3

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

Given an array nums, there is a sliding window of size k which is moving from the very left of the array to the very right. You can only see the k numbers 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,
Given nums = [1,3,-1,-3,5,3,6,7], and k = 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 assume k is always valid, ie: k is always smaller than input array's size for non-empty array.

 

This problem is similar with Data Stream Median. 

 

Solution 1. Sort the array that is inside the sliding window and get the median each time.

Runtime:  O((n - k + 1) * k * log k) 

Space: O(n)

 

Solution 2. Quick select of each window to get median.

Runtime: O((n - k + 1) * k)

 

Solution 3. Use one max and one min priority queue to dynamically maintain the medain.

The algorithm is as following.

1. Init one max pq and one min pq and add the first k numbers to them. 

2. Get the first median.

3. Add nums[i] to min/max pq and remove nums[i - k] from min/max pq. Maintain the following invariants during this step.

The max pq has 1 more element than the min pq if we've processed an odd number of elements;

They have the same number of elements if we've processed an even number of elements; 

 

public class Solution {
    /**
     * @param nums: A list of integers.
     * @return: The median of the element inside the window at each moving.
     */
    public ArrayList<Integer> medianSlidingWindow(int[] nums, int k) {
        ArrayList<Integer> medians = new ArrayList<Integer>();
        if(nums == null || nums.length == 0 || k == 0 || nums.length < k){
            return medians;
        }
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(1, Collection.reverseOrder());
        PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();
        maxHeap.add(nums[0]);
        for(int i = 1; i < k; i++){
            addNumToHeap(nums[i], maxHeap, minHeap);
        }
        medians.add(maxHeap.peek());
        for(int i = k; i < nums.length; i++){
            addNumToHeap(nums[i], maxHeap, minHeap);    
            removeNumFromHeap(nums[i - k], maxHeap, minHeap);
            medians.add(maxHeap.peek());
        }
        return medians;
    }
    private void addNumToHeap(int num, PriorityQueue<Integer> maxHeap,
                              PriorityQueue<Integer> minHeap){
        if(num <= maxHeap.peek()){
            maxHeap.add(num);
        }
        else{
            minHeap.add(num);
        }
        if(maxHeap.size() - minHeap.size() > 1){
            minHeap.add(maxHeap.poll());
        }
        if(minHeap.size() > maxHeap.size()){
            maxHeap.add(minHeap.poll());
        }
    }
    private void removeNumFromHeap(int num, PriorityQueue<Integer> maxHeap,
                                    PriorityQueue<Integer> minHeap){
        if(num <= maxHeap.peek()){
            maxHeap.remove(num);
        } 
        else{
            minHeap.remove(num);
        }
        if(maxHeap.size() - minHeap.size() > 1){
            minHeap.add(maxHeap.poll());
        }
        if(minHeap.size() > maxHeap.size()){
            maxHeap.add(minHeap.poll());
        }        
    }
}

 

Open Question: the removal of an arbitrary element in a priority queue takes O(n) time, not O(log n). n is the total number of elements in the priority queue. So even solution 3 does not achieve O(n * log n) runtime.  Is there a substitute data structure that supports O(log n) remove operation? 

One way to support O(logN) arbitrary removal is to implement your own binary heap that adds a mapping from a heap node to its index in the array. The other way to achieve the O(N*logK) runtime is to use TreeSet, instead of Priority Queue, shown as in solution 4.

 

Solution 4. Two TreeSets, O(N * logK) runtime

1. Define helper class Entry with two fields, value and its index. Define getter for sorting. No need to overwrite hashCode and equals as we are using TreeSet, not HashSet.

2. Define two TreeSet ts1 and ts2, using entry value then entry index as comparison.(This is needed because we'll have duplicated values). ts1 stores all values that are <= the current median; ts2 stores all values that are >= the current median. 

3. Add entry one by one, when there is an entry out of the current window, remove it. Both add and remove operations keep ts1 and ts2 balanced: ts1.size() is never more than ts2.size() + 1. Each balance operation takes O(logK) time.

 

 

class Solution {
    class Entry {
        int v, idx;
        Entry(int v, int idx) {
            this.v = v; 
            this.idx = idx;
        }
        int getV() {
            return v;
        }
        int getIdx() {
            return idx;
        }
    }
    private TreeSet<Entry> set1 = new TreeSet<>(Comparator.comparingInt(Entry::getV).thenComparing(Entry::getIdx));
    private TreeSet<Entry> set2 = new TreeSet<>(Comparator.comparingInt(Entry::getV).thenComparing(Entry::getIdx));
    public double[] medianSlidingWindow(int[] nums, int k) {
        double[] ans = new double[nums.length - k + 1];
        
        for(int i = 0; i < k; i++) {
            Entry e = new Entry(nums[i], i);
            addEntry(e);
        }
        ans[0] = getMedian(k);
        
        int j = 1;
        for(int i = k; i < nums.length; i++) {
            Entry rm = new Entry(nums[i - k], i - k);
            Entry ne = new Entry(nums[i], i);
            removeEntry(rm);
            addEntry(ne);
            ans[j] = getMedian(k);
            j++;
        }
        return ans;
    }
    private void addEntry(Entry e) {
        if(set1.size() == 0 || set1.last().v >= e.v) {
            set1.add(e);
            if(set1.size() > set2.size() + 1) {
                set2.add(set1.pollLast());
            }
        }
        else {
            set2.add(e);
            if(set2.size() > set1.size()) {
                set1.add(set2.pollFirst());
            }
        }        
    }
    private void removeEntry(Entry e) {
        if(set1.contains(e)) {
            set1.remove(e);
        }
        else {
            set2.remove(e);
        }
        if(set1.size() > set2.size() + 1) {
            set2.add(set1.pollLast());
        }
        else if(set2.size() > set1.size()) {
            set1.add(set2.pollFirst());
        }
    }
    private double getMedian(int k) {
        if(k % 2 != 0) {
            return (double)set1.last().v;
        }
        double sum = (double)set1.last().v + set2.first().v;
        return sum / 2;        
    }
}

 

 

 

 

 

Related Problems

Find Median From Data Stream

Sliding Window Maximum