Question

480. 滑动窗口中位数

English Version

题目描述

中位数是有序序列最中间的那个数。如果序列的长度是偶数，则没有最中间的数；此时中位数是最中间的两个数的平均数。

例如：

• [2,3,4]，中位数是 3
• [2,3]，中位数是 (2 + 3) / 2 = 2.5

给你一个数组 _nums_，有一个长度为 _k_ 的窗口从最左端滑动到最右端。窗口中有 _k_ 个数，每次窗口向右移动 _1_ 位。你的任务是找出每次窗口移动后得到的新窗口中元素的中位数，并输出由它们组成的数组。

 

示例：

给出 _nums_ = [1,3,-1,-3,5,3,6,7]，以及 _k_ = 3。

窗口位置            中位数
---------------         -----
[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

 因此，返回该滑动窗口的中位数数组 [1,-1,-1,3,5,6]。

 

提示：

• 你可以假设 k 始终有效，即：k 始终小于等于输入的非空数组的元素个数。
• 与真实值误差在 10 ^ -5 以内的答案将被视作正确答案。

解法

方法一：双优先队列（大小根堆） + 延迟删除

我们可以使用两个优先队列（大小根堆）维护当前窗口中的元素，其中一个优先队列存储当前窗口中较小的一半元素，另一个优先队列存储当前窗口中较大的一半元素。这样，当前窗口的中位数就是两个优先队列的堆顶元素的平均值或其中的一个。

我们设计一个类 MedianFinder，用于维护当前窗口中的元素。该类包含以下方法：

• add_num(num)：将 num 加入当前窗口中。
• find_median()：返回当前窗口中元素的中位数。
• remove_num(num)：将 num 从当前窗口中移除。
• prune(pq)：如果堆顶元素在延迟删除字典 delayed 中，则将其从堆顶弹出，并从该元素的延迟删除次数中减一。如果该元素的延迟删除次数为零，则将其从延迟删除字典中删除。
• rebalance()：如果较小的一半元素的数量比较大的一半元素的数量多 2 个，则将较大的一半元素的堆顶元素加入较小的一半元素中；如果较小的一半元素的数量比较大的一半元素的数量少，则将较大的一半元素的堆顶元素加入较小的一半元素中。

在 add_num(num) 方法中，我们先考虑将 num 加入较小的一半元素中，如果 num 大于较大的一半元素的堆顶元素，则将 num 加入较大的一半元素中。然后我们调用 rebalance() 方法，使得两个优先队列的大小之差不超过 $1$。

在 remove_num(num) 方法中，我们将 num 的延迟删除次数加一。然后我们将 num 与较小的一半元素的堆顶元素进行比较，如果 num 小于等于较小的一半元素的堆顶元素，则更新较小的一半元素的大小，并且调用 prune() 方法，使得较小的一半元素的堆顶元素不在延迟删除字典中。否则，我们更新较大的一半元素的大小，并且调用 prune() 方法，使得较大的一半元素的堆顶元素不在延迟删除字典中。

在 find_median() 方法中，如果当前窗口的大小为奇数，则返回较小的一半元素的堆顶元素；否则，返回较小的一半元素的堆顶元素与较大的一半元素的堆顶元素的平均值。

在 prune(pq) 方法中，如果堆顶元素在延迟删除字典中，则将其从堆顶弹出，并从该元素的延迟删除次数中减一。如果该元素的延迟删除次数为零，则将其从延迟删除字典中删除。

在 rebalance() 方法中，如果较小的一半元素的数量比较大的一半元素的数量多 2 个，则将较大的一半元素的堆顶元素加入较小的一半元素中；如果较小的一半元素的数量比较大的一半元素的数量少，则将较大的一半元素的堆顶元素加入较小的一半元素中。

时间复杂度 $O(n \times \log n)$，空间复杂度 $O(n)$。其中 $n$ 为数组 nums 的长度。

class MedianFinder:
  def __init__(self, k: int):
    self.k = k
    self.small = []
    self.large = []
    self.delayed = defaultdict(int)
    self.small_size = 0
    self.large_size = 0

  def add_num(self, num: int):
    if not self.small or num <= -self.small[0]:
      heappush(self.small, -num)
      self.small_size += 1
    else:
      heappush(self.large, num)
      self.large_size += 1
    self.rebalance()

  def find_median(self) -> float:
    return -self.small[0] if self.k & 1 else (-self.small[0] + self.large[0]) / 2

  def remove_num(self, num: int):
    self.delayed[num] += 1
    if num <= -self.small[0]:
      self.small_size -= 1
      if num == -self.small[0]:
        self.prune(self.small)
    else:
      self.large_size -= 1
      if num == self.large[0]:
        self.prune(self.large)
    self.rebalance()

  def prune(self, pq: List[int]):
    sign = -1 if pq is self.small else 1
    while pq and sign * pq[0] in self.delayed:
      self.delayed[sign * pq[0]] -= 1
      if self.delayed[sign * pq[0]] == 0:
        self.delayed.pop(sign * pq[0])
      heappop(pq)

  def rebalance(self):
    if self.small_size > self.large_size + 1:
      heappush(self.large, -heappop(self.small))
      self.small_size -= 1
      self.large_size += 1
      self.prune(self.small)
    elif self.small_size < self.large_size:
      heappush(self.small, -heappop(self.large))
      self.large_size -= 1
      self.small_size += 1
      self.prune(self.large)


class Solution:
  def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
    finder = MedianFinder(k)
    for x in nums[:k]:
      finder.add_num(x)
    ans = [finder.find_median()]
    for i in range(k, len(nums)):
      finder.add_num(nums[i])
      finder.remove_num(nums[i - k])
      ans.append(finder.find_median())
    return ans

class MedianFinder {
  private PriorityQueue<Integer> small = new PriorityQueue<>(Comparator.reverseOrder());
  private PriorityQueue<Integer> large = new PriorityQueue<>();
  private Map<Integer, Integer> delayed = new HashMap<>();
  private int smallSize;
  private int largeSize;
  private int k;

  public MedianFinder(int k) {
    this.k = k;
  }

  public void addNum(int num) {
    if (small.isEmpty() || num <= small.peek()) {
      small.offer(num);
      ++smallSize;
    } else {
      large.offer(num);
      ++largeSize;
    }
    rebalance();
  }

  public double findMedian() {
    return (k & 1) == 1 ? small.peek() : ((double) small.peek() + large.peek()) / 2;
  }

  public void removeNum(int num) {
    delayed.merge(num, 1, Integer::sum);
    if (num <= small.peek()) {
      --smallSize;
      if (num == small.peek()) {
        prune(small);
      }
    } else {
      --largeSize;
      if (num == large.peek()) {
        prune(large);
      }
    }
    rebalance();
  }

  private void prune(PriorityQueue<Integer> pq) {
    while (!pq.isEmpty() && delayed.containsKey(pq.peek())) {
      if (delayed.merge(pq.peek(), -1, Integer::sum) == 0) {
        delayed.remove(pq.peek());
      }
      pq.poll();
    }
  }

  private void rebalance() {
    if (smallSize > largeSize + 1) {
      large.offer(small.poll());
      --smallSize;
      ++largeSize;
      prune(small);
    } else if (smallSize < largeSize) {
      small.offer(large.poll());
      --largeSize;
      ++smallSize;
      prune(large);
    }
  }
}

class Solution {
  public double[] medianSlidingWindow(int[] nums, int k) {
    MedianFinder finder = new MedianFinder(k);
    for (int i = 0; i < k; ++i) {
      finder.addNum(nums[i]);
    }
    int n = nums.length;
    double[] ans = new double[n - k + 1];
    ans[0] = finder.findMedian();
    for (int i = k; i < n; ++i) {
      finder.addNum(nums[i]);
      finder.removeNum(nums[i - k]);
      ans[i - k + 1] = finder.findMedian();
    }
    return ans;
  }
}

class MedianFinder {
public:
  MedianFinder(int k) {
    this->k = k;
  }

  void addNum(int num) {
    if (small.empty() || num <= small.top()) {
      small.push(num);
      ++smallSize;
    } else {
      large.push(num);
      ++largeSize;
    }
    reblance();
  }

  void removeNum(int num) {
    ++delayed[num];
    if (num <= small.top()) {
      --smallSize;
      if (num == small.top()) {
        prune(small);
      }
    } else {
      --largeSize;
      if (num == large.top()) {
        prune(large);
      }
    }
    reblance();
  }

  double findMedian() {
    return k & 1 ? small.top() : ((double) small.top() + large.top()) / 2.0;
  }

private:
  priority_queue<int> small;
  priority_queue<int, vector<int>, greater<int>> large;
  unordered_map<int, int> delayed;
  int smallSize = 0;
  int largeSize = 0;
  int k;

  template <typename T>
  void prune(T& pq) {
    while (!pq.empty() && delayed[pq.top()]) {
      if (--delayed[pq.top()] == 0) {
        delayed.erase(pq.top());
      }
      pq.pop();
    }
  }

  void reblance() {
    if (smallSize > largeSize + 1) {
      large.push(small.top());
      small.pop();
      --smallSize;
      ++largeSize;
      prune(small);
    } else if (smallSize < largeSize) {
      small.push(large.top());
      large.pop();
      ++smallSize;
      --largeSize;
      prune(large);
    }
  }
};

class Solution {
public:
  vector<double> medianSlidingWindow(vector<int>& nums, int k) {
    MedianFinder finder(k);
    for (int i = 0; i < k; ++i) {
      finder.addNum(nums[i]);
    }
    vector<double> ans = {finder.findMedian()};
    for (int i = k; i < nums.size(); ++i) {
      finder.addNum(nums[i]);
      finder.removeNum(nums[i - k]);
      ans.push_back(finder.findMedian());
    }
    return ans;
  }
};

type MedianFinder struct {
  small        hp
  large        hp
  delayed        map[int]int
  smallSize, largeSize int
  k          int
}

func Constructor(k int) MedianFinder {
  return MedianFinder{hp{}, hp{}, map[int]int{}, 0, 0, k}
}

func (this *MedianFinder) AddNum(num int) {
  if this.small.Len() == 0 || num <= -this.small.IntSlice[0] {
    heap.Push(&this.small, -num)
    this.smallSize++
  } else {
    heap.Push(&this.large, num)
    this.largeSize++
  }
  this.rebalance()
}

func (this *MedianFinder) FindMedian() float64 {
  if this.k&1 == 1 {
    return float64(-this.small.IntSlice[0])
  }
  return float64(-this.small.IntSlice[0]+this.large.IntSlice[0]) / 2
}

func (this *MedianFinder) removeNum(num int) {
  this.delayed[num]++
  if num <= -this.small.IntSlice[0] {
    this.smallSize--
    if num == -this.small.IntSlice[0] {
      this.prune(&this.small)
    }
  } else {
    this.largeSize--
    if num == this.large.IntSlice[0] {
      this.prune(&this.large)
    }
  }
  this.rebalance()
}

func (this *MedianFinder) prune(pq *hp) {
  sign := 1
  if pq == &this.small {
    sign = -1
  }
  for pq.Len() > 0 && this.delayed[sign*pq.IntSlice[0]] > 0 {
    this.delayed[sign*pq.IntSlice[0]]--
    if this.delayed[sign*pq.IntSlice[0]] == 0 {
      delete(this.delayed, sign*pq.IntSlice[0])
    }
    heap.Pop(pq)
  }
}

func (this *MedianFinder) rebalance() {
  if this.smallSize > this.largeSize+1 {
    heap.Push(&this.large, -heap.Pop(&this.small).(int))
    this.smallSize--
    this.largeSize++
    this.prune(&this.small)
  } else if this.smallSize < this.largeSize {
    heap.Push(&this.small, -heap.Pop(&this.large).(int))
    this.smallSize++
    this.largeSize--
    this.prune(&this.large)
  }
}

func medianSlidingWindow(nums []int, k int) []float64 {
  finder := Constructor(k)
  for _, num := range nums[:k] {
    finder.AddNum(num)
  }
  ans := []float64{finder.FindMedian()}
  for i := k; i < len(nums); i++ {
    finder.AddNum(nums[i])
    finder.removeNum(nums[i-k])
    ans = append(ans, finder.FindMedian())
  }
  return ans
}

type hp struct{ sort.IntSlice }

func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] }
func (h *hp) Push(v any)    { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
  a := h.IntSlice
  v := a[len(a)-1]
  h.IntSlice = a[:len(a)-1]
  return v
}