# Java using two Tree Sets - O(n logk)

• Inspired by this solution. to the problem: 295. Find Median from Data Stream

However instead of using two priority queue's we use two `Tree Sets` as we want `O(logk)` for `remove(element)`. Priority Queue would have been `O(k)` for `remove(element)` giving us an overall time complexity of `O(nk)` instead of `O(nlogk)`.

However there is an issue when it comes to duplicate values so to get around this we store the `index` into `nums` in our `Tree Set`. To break ties in our Tree Set comparator we compare the index.

``````public double[] medianSlidingWindow(int[] nums, int k) {
Comparator<Integer> comparator = (a, b) -> nums[a] != nums[b] ? Integer.compare(nums[a], nums[b]) : a - b;
TreeSet<Integer> left = new TreeSet<>(comparator.reversed());
TreeSet<Integer> right = new TreeSet<>(comparator);

Supplier<Double> median = (k % 2 == 0) ?
() -> ((double) nums[left.first()] + nums[right.first()]) / 2 :
() -> (double) nums[right.first()];

// balance lefts size and rights size (if not equal then right will be larger by one)
Runnable balance = () -> { while (left.size() > right.size()) right.add(left.pollFirst()); };

double[] result = new double[nums.length - k + 1];

for (int i = 0; i < k; i++) left.add(i);
balance.run(); result[0] = median.get();

for (int i = k, r = 1; i < nums.length; i++, r++) {
// remove tail of window from either left or right
if(!left.remove(i - k)) right.remove(i - k);

// add next num, this will always increase left size

// rebalance left and right, then get median from them
balance.run(); result[r] = median.get();
}

return result;
}
``````

• I have learned a lot from this post. Thank you very much.

• Is `remove(element)` really O(1) for `TreeSet`?

The Java API for TreeSet says,
`This implementation provides guaranteed log(n) time cost for the basic operations (add, remove and contains).`

• @ppeterson

My bad! Good catch, it should be O(logk) for remove (will update the post). It doesn't change the overall runtime complexity. If we would have used PriorityQueue then remove(object) would be linear time instead of logarithmic, affecting the overall runtime complexity.

• This post is deleted!

• Definitely awesome solution. Thanks!

• java lambdas are magical

• Great idea! I re-write it using another way:

``````class Solution {
static class myInteger{
int val;
int index;
myInteger(int val,int index){
this.val = val;
this.index = index;
}
}
public double[] medianSlidingWindow(int[] nums, int k) {
TreeSet<myInteger> minheap = new TreeSet<>(new Comparator<myInteger>(){
public int compare(myInteger a,myInteger b){
if(a.val!=b.val){
if(a.val<b.val){
return -1;
}else{
return 1;
}
}else{
return a.index-b.index;
}
}
});
TreeSet<myInteger> maxheap = new TreeSet<>(new Comparator<myInteger>(){
public int compare(myInteger a,myInteger b){
if(a.val!=b.val){
if(a.val<b.val){
return -1;
}else{
return 1;
}
}else{
return a.index - b.index;
}
}
});

Deque<myInteger> deque  = new ArrayDeque<>();
double[] res = new double[nums.length-k+1];
for(int i=0;i<k;i++){
myInteger temp = new myInteger(nums[i],i);
deque.offer(temp);
}
balance(minheap,maxheap);
//System.out.println("size of minheap :" + minheap.size() + " and the size of maxheap is : " + maxheap.size());

res[0] = getmedian(minheap,maxheap);
int p=1;
for(int i=k;i<nums.length;i++){
myInteger removeEle = deque.pollFirst();
if(minheap.contains(removeEle)){
minheap.remove(removeEle);
}else{
maxheap.remove(removeEle);
}
myInteger newEle = new myInteger(nums[i],i);
deque.offer(newEle);
balance(minheap,maxheap);
res[p++] = getmedian(minheap,maxheap);
}
return res;

}

public double getmedian(TreeSet<myInteger> minHeap,TreeSet<myInteger> maxHeap){
if(minHeap.size()>maxHeap.size()){
return (double)minHeap.first().val;
}
return ((double)minHeap.first().val+(double)maxHeap.last().val)/2.0;
}

public void balance(TreeSet<myInteger> minHeap,TreeSet<myInteger> maxHeap){
while(maxHeap.size()>minHeap.size()){