LeetCode 480. 滑动窗口中位数

题目描述

480. 滑动窗口中位数

思路分析

解法一:双堆 + 延迟删除(推荐)

核心思路

  • 用一个大根堆维护较小一半元素,小根堆维护较大一半元素。
  • 通过延迟删除哈希表记录需移除的元素,堆顶弹出时再清理。
  • 保持两堆大小平衡,随时可取到窗口中位数。


复杂度分析

  • 时间复杂度:O(n log k),其中 n 表示数组长度,k 表示窗口大小。
  • 空间复杂度:O(k),堆与延迟删除表占用线性空间。
class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        DualHeap dh = new DualHeap(k);

        for (int i = 0; i < k; i++) {
            dh.insert(nums[i]);
        }

        double[] res = new double[nums.length - k + 1];
        res[0] = dh.getMedian();

        for (int i = k; i < nums.length; i++) {
            dh.insert(nums[i]);
            dh.erase(nums[i - k]);
            res[i - k + 1] = dh.getMedian();
        }

        return res;
    }

    private static class DualHeap {
        private final PriorityQueue<Integer> small;
        private final PriorityQueue<Integer> large;
        private final Map<Integer, Integer> delayed;
        private final int k;
        private int smallSize;
        private int largeSize;

        DualHeap(int k) {
            this.k = k;
            this.small = new PriorityQueue<>((a, b) -> b - a);
            this.large = new PriorityQueue<>();
            this.delayed = new HashMap<>();
        }

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

        public void erase(int num) {
            delayed.put(num, delayed.getOrDefault(num, 0) + 1);

            if (!small.isEmpty() && num <= small.peek()) {
                smallSize--;
                if (num == small.peek()) {
                    prune(small);
                }
            } else {
                largeSize--;
                if (!large.isEmpty() && num == large.peek()) {
                    prune(large);
                }
            }

            makeBalance();
        }

        public double getMedian() {
            if ((k & 1) == 1) {
                return small.peek();
            }
            return ((double) small.peek() + large.peek()) / 2.0;
        }

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

        private void prune(PriorityQueue<Integer> heap) {
            while (!heap.isEmpty()) {
                int num = heap.peek();
                Integer cnt = delayed.get(num);
                if (cnt == null) {
                    break;
                }
                if (cnt == 1) {
                    delayed.remove(num);
                } else {
                    delayed.put(num, cnt - 1);
                }
                heap.poll();
            }
        }
    }
}
func medianSlidingWindow(nums []int, k int) []float64 {
	dh := NewDualHeap(k)
	for i := 0; i < k; i++ {
		dh.Insert(nums[i])
	}

	res := make([]float64, len(nums)-k+1)
	res[0] = dh.Median()

	for i := k; i < len(nums); i++ {
		dh.Insert(nums[i])
		dh.Erase(nums[i-k])
		res[i-k+1] = dh.Median()
	}

	return res
}

type DualHeap struct {
	small     *IntHeap
	large     *IntHeap
	delayed   map[int]int
	k         int
	smallSize int
	largeSize int
}

func NewDualHeap(k int) *DualHeap {
	small := &IntHeap{isMin: false}
	large := &IntHeap{isMin: true}
	heap.Init(small)
	heap.Init(large)

	return &DualHeap{
		small:   small,
		large:   large,
		delayed: make(map[int]int),
		k:       k,
	}
}

func (dh *DualHeap) Insert(num int) {
	if dh.small.Len() == 0 || num <= dh.small.Top() {
		heap.Push(dh.small, num)
		dh.smallSize++
	} else {
		heap.Push(dh.large, num)
		dh.largeSize++
	}
	dh.makeBalance()
}

func (dh *DualHeap) Erase(num int) {
	dh.delayed[num]++

	if dh.small.Len() > 0 && num <= dh.small.Top() {
		dh.smallSize--
		if num == dh.small.Top() {
			dh.prune(dh.small)
		}
	} else {
		dh.largeSize--
		if dh.large.Len() > 0 && num == dh.large.Top() {
			dh.prune(dh.large)
		}
	}

	dh.makeBalance()
}

func (dh *DualHeap) Median() float64 {
	if dh.k%2 == 1 {
		return float64(dh.small.Top())
	}
	return (float64(dh.small.Top()) + float64(dh.large.Top())) / 2.0
}

func (dh *DualHeap) makeBalance() {
	if dh.smallSize > dh.largeSize+1 {
		heap.Push(dh.large, heap.Pop(dh.small))
		dh.smallSize--
		dh.largeSize++
		dh.prune(dh.small)
	} else if dh.smallSize < dh.largeSize {
		heap.Push(dh.small, heap.Pop(dh.large))
		dh.largeSize--
		dh.smallSize++
		dh.prune(dh.large)
	}
}

func (dh *DualHeap) prune(h *IntHeap) {
	for h.Len() > 0 {
		num := h.Top()
		cnt, ok := dh.delayed[num]
		if !ok {
			break
		}
		if cnt == 1 {
			delete(dh.delayed, num)
		} else {
			dh.delayed[num] = cnt - 1
		}
		heap.Pop(h)
	}
}

type IntHeap struct {
	data  []int
	isMin bool
}

func (h IntHeap) Len() int { return len(h.data) }
func (h IntHeap) Less(i, j int) bool {
	if h.isMin {
		return h.data[i] < h.data[j]
	}
	return h.data[i] > h.data[j]
}
func (h IntHeap) Swap(i, j int) { h.data[i], h.data[j] = h.data[j], h.data[i] }
func (h *IntHeap) Push(x any)   { h.data = append(h.data, x.(int)) }
func (h *IntHeap) Pop() any {
	old := h.data
	n := len(old)
	val := old[n-1]
	h.data = old[:n-1]
	return val
}
func (h IntHeap) Top() int { return h.data[0] }

相似题目

题目 难度 考察点
295. 数据流的中位数 困难 双堆
239. 滑动窗口最大值 困难 单调队列
703. 数据流中的第 K 大元素 简单
480. 滑动窗口中位数 困难 双堆
346. 数据流中的移动平均值 简单 队列
本文作者:
本文链接: https://hgnulb.github.io/blog/2026/80498442
版权声明: 本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议,转载请注明出处!