利口 480 滑动窗口中位数

题意：给定一个数组和一个滑动窗口，返回每一个滑动窗口的中位数

思路：遍历数组

1. 维护一个最小的treeset，和一个最大的treeset，记录当前节点的index
2. 每次根据奇偶返回当前的中位数

思想：treeset

复杂度：时间O（nlgk)，空间O(k)

class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        TreeSet<Integer> set1 = new TreeSet(new Comparator<Integer>(){
           public int compare(Integer n1, Integer n2) {
               if(nums[n2] == nums[n1])
                   return n1 - n2;
               return Integer.compare(nums[n1], nums[n2]);
           } 
        });
        TreeSet<Integer> set2 = new TreeSet(new Comparator<Integer>(){
           public int compare(Integer n1, Integer n2) {
               if(nums[n2] == nums[n1])
                   return n1 - n2;
               return Integer.compare(nums[n2], nums[n1]);
           } 
        });
        
        double[] res = new double[nums.length - k + 1];
        
        for(int i=0;i<nums.length;i++) {
            if(i>=k) {
                res[i - k] = get(set1, set2, nums);
                remove(i-k, set1, set2, nums);
            }
            add(i, set1, set2, nums);
        }
        res[nums.length - k] = get(set1, set2, nums);
        return res;
    }
    
    void add(int n1, TreeSet<Integer> set1, TreeSet<Integer> set2, int[] nums) {
        if(set1.isEmpty() || nums[n1] >= nums[set1.first()]) {
            set1.add(n1);
        } else {
            set2.add(n1);
        }
        if(set1.size() < set2.size())
            set1.add(set2.pollFirst());
        
        if(set1.size() - set2.size() > 1)
            set2.add(set1.pollFirst());
    }
        
    void remove(int n1, TreeSet<Integer> set1, TreeSet<Integer> set2, int[] nums) {
        if(set1.contains(n1)) {
            set1.remove(n1);
        } else {
            set2.remove(n1);
        }
        if(set1.size() < set2.size())
            set1.add(set2.pollFirst());
        
        if(set1.size() - set2.size() > 1)
            set2.add(set1.pollFirst());
    }
    
    double get(TreeSet<Integer> set1, TreeSet<Integer> set2, int[] nums) {
        if(set1.size() == set2.size()) {
            return ((double)nums[set1.first()] + (double)nums[set2.first()]) / 2;
        } else {
            return (double)nums[set1.first()];
        }
    }
}