Sliding Window Median

This page explains Java solution to problem Sliding Window Median using TreeSet data structure.

Problem Statement

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

Examples:
[2,3,4], the median is 3
[2,3], the median is (2 + 3) / 2 = 2.5

Given an array nums, there is a sliding window of size k which is moving from the very left of the array to the very right. You can only see the k numbers in the window. Each time the sliding window moves right by one position. Your job is to output the median array for each window in the original array.

Input: nums = [1,3,-1,-3,5,3,6,7], and k = 3
Output: [1,-1,-1,3,5,6]
Explanation:
Window position            Median
[1 3 -1] -3 5 3 6 7            1
1 [3 -1 -3] 5 3 6 7           -1
1 3 [-1 -3 5] 3 6 7           -1
1 3 -1 [-3 5 3] 6 7            3
1 3 -1 -3 [5 3 6] 7            5
1 3 -1 -3 5 [3 6 7]            6

Solution

package com.vc.hard;

import java.util.*;

class SlidingWindowMedian {
    public double[] medianSlidingWindow(int[] nums, int k) {
        if(nums == null) return null;
        if(nums.length == 0) return new double[0];

        TreeSet<Integer> left = new TreeSet<>(new Comparator<Integer>(){
            public int compare(Integer i1, Integer i2) {
                int cmp = Integer.compare(nums[i1], nums[i2]);
                return cmp == 0 ? Integer.compare(i1, i2) : cmp;
            }
        });

        TreeSet<Integer> right = new TreeSet<>(new Comparator<Integer>(){
            public int compare(Integer i1, Integer i2) {
                int cmp = Integer.compare(nums[i2], nums[i1]);
                return cmp == 0 ? Integer.compare(i1, i2): cmp;
            }
        });

        int n = nums.length;
        for(int i = 0; i < k; i++) {
            add(left, right, i);
        }

        double[] res = new double[n - k + 1];
        for(int i = k; i < n; i++) {
            res[i - k] = getMedian(left, right, nums);
            left.remove(i - k);
            right.remove(i - k);
            add(left, right, i);
        }

        res[res.length - 1] = getMedian(left, right, nums);
        return res;
    }

    private void add(TreeSet<Integer> left, TreeSet<Integer> right, int index) {
        if(left.size() > right.size()) {
            left.add(index);
            right.add(left.pollFirst());
        }
        else {
            right.add(index);
            left.add(right.pollFirst());
        }
    }

    private double getMedian(TreeSet<Integer> left, TreeSet<Integer> right, int[] nums) {
        if(left.size() > right.size()) {
            return nums[left.first()];
        }
        else {
            double d = (double)nums[left.first()] + (double)nums[right.first()];
            return d / 2.0;
        }
    }
}

Time Complexity

O(N * log K) Where
N is total number of elements in an input array
K is size of sliding window

Space Complexity

O(N) Where
N is total number of elements int an input array