Method 1.2 Two TreeSets

Note: How do we design delete(value) API, when we really would like to decide which one to delete

  • Node{Value, sequence #} (assume we dont know the sequence?)

  • Choice 1: remove the largest node with value 2 ==> treeSet.floor(new Node(2, globalCounter));

  • Choice 2: remove the smallest node with value 2 ==> treeSet.ceiling(new Node(2, 0))

However, since this question is a sliding window, we can get the index

class Solution {
    private int windowSize;
    private TreeSet<Element> smaller;
    private TreeSet<Element> larger;
    public double[] medianSlidingWindow(int[] nums, int k) {
        windowSize = nums.length;
        smaller = new TreeSet<>();
        larger = new TreeSet<>();
        double[] result = new double[nums.length - k + 1];
        for (int i = 0; i < nums.length; i++) {
            if (i >= k) {
                delete(nums[i - k], i - k);
            }
            add(nums[i], i);
            if (i >= k - 1) {
                result[i - k + 1] = findMedian();
            }
        }
        return result;
    }
    private void add(int num, int freq) {
        Element e = new Element(num, freq);
        if (smaller.isEmpty() || e.compareTo(smaller.last()) <= 0) {
            smaller.add(e);
        }
        else {
            larger.add(e);
        }
        rebalance();
    }
    private void delete(int num, int freq) {
        Element e = new Element(num, freq);
        if (smaller.contains(e)) {
            smaller.remove(e);
        }
        else {
            larger.remove(e);
        }
        rebalance();
    }
    private void rebalance() {
        while (smaller.size() > larger.size() + 1) {
            larger.add(smaller.pollLast());
        }
        while (larger.size() > smaller.size()){
            smaller.add(larger.pollFirst());
        }

    }
    private double findMedian() {
        if (this.smaller.size() == this.larger.size()) {
            return (double) this.smaller.last().value /2.0  + this.larger.first().value/ 2.0;
        }
        else {
            return (double) this.smaller.last().value;
        }
    }
    public class Element implements Comparable<Element> {
        int value;
        int index;
        public Element(int value, int index) {
            this.value = value;
            this.index = index;
        }
        @Override
        public int compareTo(Element that) {
            if (this.value != that.value) {
                return Integer.compare(this.value, that.value);
            }
            else {
                return Integer.compare(this.index, that.index);
            }
        }
    }
}
```

Last updated