Comparator<Node> comparator = (a, b) -> a.val != b.val ? a.val - b.val : a.index - b.index;
TreeSet<Node> maxHeap = new TreeSet<>(comparator.reversed());
TreeSet<Node> minHeap = new TreeSet<>(comparator);
public class Solution {
/**
* @param nums: A list of integers
* @param k: An integer
* @return: The median of the element inside the window at each moving
*/
public List<Integer> medianSlidingWindow(int[] nums, int k) {
List<Integer> ans = new ArrayList<Integer>();
if (nums == null || nums.length == 0) {
return ans;
}
Comparator<Node> comparator = (a,b)-> (a.val != b.val?a.val - b.val : a.index - b.index);
TreeSet<Node> maxTreeSet = new TreeSet<Node>(comparator.reversed());
TreeSet<Node> minTreeSet = new TreeSet<Node>(comparator);
int size = (k + 1) / 2;
for (int i = 0; i < k - 1; i++) {
add(maxTreeSet, minTreeSet, size, new Node(i, nums[i]));
}
for (int i = k - 1; i < nums.length; i++) {
add(maxTreeSet, minTreeSet, size, new Node(i, nums[i]));
ans.add(maxTreeSet.first().val);//median
removeNode(maxTreeSet, minTreeSet, new Node(i-k+1, nums[i-k+1]));
}
return ans;
}
private void add( TreeSet<Node> maxTreeSet, TreeSet<Node> minTreeSet, int size, Node nd) {
if (maxTreeSet.size() < size) {
maxTreeSet.add(nd);
}else{
minTreeSet.add(nd);
}
if (maxTreeSet.size() == size) {
if (minTreeSet.size() != 0 && maxTreeSet.first().val > minTreeSet.first().val) {
Node big = maxTreeSet.first();
Node small = minTreeSet.first();
maxTreeSet.remove(big);
minTreeSet.remove(small);
maxTreeSet.add(small);
minTreeSet.add(big);
}
}
}
private void removeNode( TreeSet<Node> maxTreeSet, TreeSet<Node> minTreeSet, Node nd) {
if (maxTreeSet.contains(nd)) {
maxTreeSet.remove(nd);
}else{
minTreeSet.remove(nd);
}
}
class Node {
int index;
int val;
private Node(int idx, int value) {
index = idx;
val = value;
}
}
}