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.

 

开始写了用一个BST的版本，在中值重复的情况下，要考虑在第几个数。代码非常繁琐不清晰。这题不能像295. Find Median from Data Stream那样用两个优先队列，因为包含删除操作，优先队列的删除操作是O（n），看讨论区有人用了两个BST，非常简洁。BST和优先队列比能在O（log n）时间做很多操作，优先队列最大的优势是能在O（1）时间找到最值。

两个treemap一个维护存放小于等于中值的数，一个存放大于中值的数。时间复杂度O（n log k）.

class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        TreeMap<Integer,Integer> m1=new TreeMap<>();
        TreeMap<Integer,Integer> m2=new TreeMap<>();
        for(int i=0;i<k;i++){
            if(i%2==0){
                m2.put(nums[i],m2.getOrDefault(nums[i],0)+1);
                int n=m2.firstKey();
                if(m2.get(n)==1) m2.remove(n);
                else m2.put(n,m2.get(n)-1);
                m1.put(n,m1.getOrDefault(n,0)+1);
            }else{
                m1.put(nums[i],m1.getOrDefault(nums[i],0)+1);
                int n=m1.lastKey();
                if(m1.get(n)==1) m1.remove(n);
                else m1.put(n,m1.get(n)-1);
                m2.put(n,m2.getOrDefault(n,0)+1);
            }
        }
        
        double[] res=new double[nums.length-k+1];
        if(k%2==1) res[0]=m1.lastKey();
        else res[0]=((double)m1.lastKey()+(double)m2.firstKey())/2;
        for(int i=0;i+k<nums.length;i++){
            //insert
            m1.put(nums[i+k],m1.getOrDefault(nums[i+k],0)+1);
            int l=m1.lastKey();
            if(m1.get(l)==1){
                m1.remove(l);
            }else{
                m1.put(l,m1.get(l)-1);
            }
            m2.put(l,m2.getOrDefault(l,0)+1);
            
            //delete
            if(m1.containsKey(nums[i])){
                if(m1.get(nums[i])==1){
                    m1.remove(nums[i]);
                }else{
                    m1.put(nums[i],m1.get(nums[i])-1);
                }
                int n=m2.firstKey();
                if(m2.get(n)==1){
                    m2.remove(n);
                }else{
                    m2.put(n,m2.get(n)-1);
                }
                m1.put(n,m1.getOrDefault(n,0)+1);
            }else{
                if(m2.get(nums[i])==1){
                    m2.remove(nums[i]);
                }else{
                    m2.put(nums[i],m2.get(nums[i])-1);
                }
            }
            
            //result
            if(k%2==1) res[i+1]=m1.lastKey();
            else res[i+1]=((double)m1.lastKey()+(double)m2.firstKey())/2;
        }
        return res;
    }
}