2020-09-14

class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        double[] result = new double[nums.length - k + 1];
        if (k == 1) {
            for (int i = 0; i < result.length; i ++) {
                result[i] = (double) nums[i];
            }
            return result;
        }
        // PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> (o2 < 0 && o1 > 0 ? -1 : o2 > 0 && o1 < 0 ? 1 : o2 - o1));
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        if ((k & 1) == 1) {
            for (int i = 0; i < k; i ++) {
                maxHeap.add(nums[i]);
            }
            while (maxHeap.size() > minHeap.size() + 1) {
                minHeap.add(maxHeap.poll());
            }
            result[0] = (double) maxHeap.peek();
            for (int i = k; i < nums.length; i ++) {
                removeAndPush(nums, k, maxHeap, minHeap, i);
                if (maxHeap.size() < minHeap.size()) {
                    maxHeap.add(minHeap.poll());
                }
                else if (maxHeap.size() > minHeap.size() + 1) {
                    minHeap.add(maxHeap.poll());
                }
                result[i-k+1] = (double) maxHeap.peek();
            }
        }
        else {
            for (int i = 0; i < k; i ++) {
                maxHeap.add(nums[i]);
            }
            while (maxHeap.size() > minHeap.size()) {
                minHeap.add(maxHeap.poll());
            }
            result[0] = ((double)maxHeap.peek() + (double)minHeap.peek()) / 2;
            for (int i = k; i < nums.length; i ++) {
                removeAndPush(nums, k, maxHeap, minHeap, i);
                if (maxHeap.size() > minHeap.size()) {
                    minHeap.add(maxHeap.poll());
                }
                else if (minHeap.size() > maxHeap.size()) {
                    maxHeap.add(minHeap.poll());
                }
                // System.out.println("------");
                // System.out.println(maxHeap);
                // System.out.println(minHeap);
                result[i-k+1] = ((double)maxHeap.peek() + (double)minHeap.peek()) / 2;
            }
        }
        return result;
    }

    private void removeAndPush(int[] nums, int k, PriorityQueue<Integer> maxHeap, PriorityQueue<Integer> minHeap, int i) {
        if (maxHeap.size()!= 0 && nums[i - k] <= maxHeap.peek()) {
            maxHeap.remove(nums[i - k]);
        } else {
            minHeap.remove(nums[i - k]);
        }
        if (maxHeap.size() == 0) {
            minHeap.add(nums[i]);
            maxHeap.add(minHeap.poll());
        }
        else if (nums[i] <= maxHeap.peek()) {
            maxHeap.add(nums[i]);
        } else {
            minHeap.add(nums[i]);
        }
    }

    // public static void main(String[] args) {
    //     Solution s = new Solution();
    //     s.medianSlidingWindow(new int[] {1,1,1,1,7,1,3,-3,5,3,6,7,1,3,-1,-3,5,3,6,7,1,3,-1,-3,5,3,6,7}, 4);
    // }
}