Golang | Leetcode Golang题解之第480题滑动窗口中位数

题目：

题解：

type hp struct {
    sort.IntSlice
    size int
}
func (h *hp) Push(v interface{}) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() interface{}   { a := h.IntSlice; v := a[len(a)-1]; h.IntSlice = a[:len(a)-1]; return v }
func (h *hp) push(v int)         { h.size++; heap.Push(h, v) }
func (h *hp) pop() int           { h.size--; return heap.Pop(h).(int) }
func (h *hp) prune() {
    for h.Len() > 0 {
        num := h.IntSlice[0]
        if h == small {
            num = -num
        }
        if d, has := delayed[num]; has {
            if d > 1 {
                delayed[num]--
            } else {
                delete(delayed, num)
            }
            heap.Pop(h)
        } else {
            break
        }
    }
}

var delayed map[int]int
var small, large *hp

func medianSlidingWindow(nums []int, k int) []float64 {
    delayed = map[int]int{} // 哈希表，记录「延迟删除」的元素，key 为元素，value 为需要删除的次数
    small = &hp{}           // 大根堆，维护较小的一半元素
    large = &hp{}           // 小根堆，维护较大的一半元素
    makeBalance := func() {
        // 调整 small 和 large 中的元素个数，使得二者的元素个数满足要求
        if small.size > large.size+1 { // small 比 large 元素多 2 个
            large.push(-small.pop())
            small.prune() // small 堆顶元素被移除，需要进行 prune
        } else if small.size < large.size { // large 比 small 元素多 1 个
            small.push(-large.pop())
            large.prune() // large 堆顶元素被移除，需要进行 prune
        }
    }
    insert := func(num int) {
        if small.Len() == 0 || num <= -small.IntSlice[0] {
            small.push(-num)
        } else {
            large.push(num)
        }
        makeBalance()
    }
    erase := func(num int) {
        delayed[num]++
        if num <= -small.IntSlice[0] {
            small.size--
            if num == -small.IntSlice[0] {
                small.prune()
            }
        } else {
            large.size--
            if num == large.IntSlice[0] {
                large.prune()
            }
        }
        makeBalance()
    }
    getMedian := func() float64 {
        if k&1 > 0 {
            return float64(-small.IntSlice[0])
        }
        return float64(-small.IntSlice[0]+large.IntSlice[0]) / 2
    }

    for _, num := range nums[:k] {
        insert(num)
    }
    n := len(nums)
    ans := make([]float64, 1, n-k+1)
    ans[0] = getMedian()
    for i := k; i < n; i++ {
        insert(nums[i])
        erase(nums[i-k])
        ans = append(ans, getMedian())
    }
    return ans
}