g0401_0500.s0480_sliding_window_median.Solution Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of leetcode-in-java21 Show documentation
Java-based LeetCode algorithm problem solutions, regularly updated
There is a newer version: 1.38
package g0401_0500.s0480_sliding_window_median;

// #Hard #Array #Hash_Table #Heap_Priority_Queue #Sliding_Window
// #2022_07_20_Time_60_ms_(70.77%)_Space_56.3_MB_(17.95%)

import java.util.Comparator;
import java.util.TreeSet;

/**
 * 480 - Sliding Window Median\.
 *
 * Hard
 *
 * The **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 values.
 *
 * *   For examples, if `arr = [2,3,4]`, the median is `3`.
 * *   For examples, if `arr = [1,2,3,4]`, the median is `(2 + 3) / 2 = 2.5`.
 *
 * You are given an integer array `nums` and an integer `k`. 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.
 *
 * Return _the median array for each window in the original array_. Answers within 10^-5 of the actual value will be accepted.
 *
 * **Example 1:**
 *
 * **Input:** nums = [1,3,-1,-3,5,3,6,7], k = 3
 *
 * **Output:** [1.00000,-1.00000,-1.00000,3.00000,5.00000,6.00000]
 *
 * **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
 *
 * **Example 2:**
 *
 * **Input:** nums = [1,2,3,4,2,3,1,4,2], k = 3
 *
 * **Output:** [2.00000,3.00000,3.00000,3.00000,2.00000,3.00000,2.00000]
 *
 * **Constraints:**
 *
 * *   1 <= k <= nums.length <= 10⁵
 * *   -2³¹ <= nums[i] <= 2³¹ - 1
**/
@SuppressWarnings("java:S3012")
public class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        if (nums == null || k <= 0) {
            throw new IllegalArgumentException("Input is invalid");
        }
        int len = nums.length;
        double[] result = new double[len - k + 1];
        if (k == 1) {
            for (int i = 0; i < len; i++) {
                result[i] = nums[i];
            }
            return result;
        }
        Comparator comparator =
                (a, b) ->
                        (nums[a] != nums[b]
                                ? Integer.compare(nums[a], nums[b])
                                : Integer.compare(a, b));
        TreeSet smallNums = new TreeSet<>(comparator.reversed());
        TreeSet largeNums = new TreeSet<>(comparator);
        for (int i = 0; i < len; i++) {
            if (i >= k) {
                removeElement(smallNums, largeNums, i - k);
            }
            addElement(smallNums, largeNums, i);
            if (i >= k - 1) {
                result[i - (k - 1)] = getMedian(smallNums, largeNums, nums);
            }
        }
        return result;
    }

    private void addElement(TreeSet smallNums, TreeSet largeNums, int idx) {
        smallNums.add(idx);
        largeNums.add(smallNums.pollFirst());
        if (smallNums.size() < largeNums.size()) {
            smallNums.add(largeNums.pollFirst());
        }
    }

    private void removeElement(TreeSet smallNums, TreeSet largeNums, int idx) {
        if (largeNums.contains(idx)) {
            largeNums.remove(idx);
            if (smallNums.size() == largeNums.size() + 2) {
                largeNums.add(smallNums.pollFirst());
            }
        } else {
            smallNums.remove(idx);
            if (smallNums.size() < largeNums.size()) {
                smallNums.add(largeNums.pollFirst());
            }
        }
    }

    private double getMedian(TreeSet smallNums, TreeSet largeNums, int[] nums) {
        if (smallNums.size() == largeNums.size()) {
            return ((double) nums[smallNums.first()] + nums[largeNums.first()]) / 2;
        }
        return nums[smallNums.first()];
    }
}