medianSlidingWindow.java

class Solution {
    private TreeMap<Integer, Integer> map;
    private int index;
    private int middle;
    private int k;
    
    public double[] medianSlidingWindow(int[] nums, int k) {
        int m = (k + 1) % 2;
        this.k = k;
        
        map = new TreeMap<Integer, Integer>();

        ArrayList<Integer> list = new ArrayList<Integer>();
        ArrayList<Double> ret = new ArrayList<Double>();

        for (int i = 0; i < k; i++) {
            push(nums[i]);
            list.add(nums[i]);
        }
        
        Collections.sort(list);
        
        int center = (k - 1) / 2;
        middle = list.get(center);
        index = 0;
        for (int i = center - 1; i >= 0; i--) {
            if (list.get(i) != middle) {
                break;
            }
            index++;
        }
        
        for (int i = 0; (i + k) <= nums.length; i++) {
            double sum = (double) middle;
            if (m != 0) {
                sum += (double) next();
                sum /= 2;
            }
            
            ret.add(sum);
            
            if (i + k != nums.length) {
                add(nums[i + k]);
                remove(nums[i]);
            }
        }

        return ret.stream().mapToDouble(Double::doubleValue).toArray(); 
    }

    private void remove(int num) {
        if (num == middle) {
            int count = map.get(middle) != null ? map.get(middle) : 0;
            if ((k % 2 != 0) && (index + 1 == count)) {
                index = 0;
                middle = map.higherKey(middle) == null ? map.lastKey() : map.higherKey(middle);
            } else if (k % 2 == 0) {
                if (index == 0) {
                    middle = map.lowerKey(middle);
                    index = map.get(middle) - 1;
                } else {
                    index--;
                }
            }
            pop(num);
        } else {
            pop(num);
            if ((k % 2 == 0) && (num > middle)) {
                if (index > 0) {
                    index--;
                } else {
                    middle = map.lowerKey(middle);
                    index = map.get(middle) - 1;
                }
            } else if ((k % 2 != 0) && (num < middle)) {
                int count = map.get(middle) != null ? map.get(middle) : 0;
                if (index + 1 < count) {
                    index++;
                } else {
                    index = 0;
                    middle = map.higherKey(middle);
                }
            }
        }
    }

    private void add(int num) {
        push(num);
        if ((k % 2 == 0) && (num >= middle)) {
            int count = map.get(middle) != null ? map.get(middle) : 0;
            if (index + 1 < count) {
                index++;
            } else {
                index = 0;
                middle = map.higherKey(middle);
            }
        } else if ((k % 2 != 0) && (num < middle)){
            if (index > 0) {
                index--;
            } else {
                middle = map.lowerKey(middle);
                index = map.get(middle) - 1;
            }
        }
    }
    
    private int next() {
        int count = map.get(middle) != null ? map.get(middle) : 0;
        if (index + 1 < count) {
            return middle;
        } else {
            return map.higherKey(middle);
        }
    }
    
    private void push(int num) {
        int count = map.get(num) != null ? map.get(num) : 0;
        map.put(num, count + 1);
    }

    private void pop(int num) {
        int count = map.get(num) != null ? map.get(num) : 0;
        if (count == 1) {
            map.remove(num);
        } else {
            map.put(num, count - 1);            
        }
    }
}