480. 滑动窗口中位数

241 阅读1分钟

题目: 算法:

方法一:暴力
TLE

func medianSlidingWindow(nums []int, k int) []float64 {
	ans := make([]float64, 0)
	tmp := make([]int, k)
	for i := 0; i+k <= len(nums); i++ {
		for j := 0; j < k; j++ {
			tmp[j] = nums[j+i]
		}

		sort.Ints(tmp)
		mid := k / 2
		if k%2 == 0 {
			ans = append(ans, float64(tmp[mid]+tmp[mid-1])/2)
		} else {
			ans = append(ans, float64(tmp[mid]))

		}
	}
	return ans
}

方法二:大顶堆+小顶堆
我们有什么方法快速查找中位数?大顶堆+小顶堆

func medianSlidingWindow(nums []int, k int) []float64 {
    ans := make([]float64, 0)
    // small是大顶堆, big是小顶堆
    big := &SmallHeap{}
    small := &BigHeap{}
    heap.Init(big)
    heap.Init(small)
    
    
    // 初始化大顶堆和小顶堆
    // 如果i < k, 初始化,全部加入small,
    // i == k 从small分一半到big, 分完之后,small.len() >= big.len()
    for i := 0 ;i < k; i ++ {
        heap.Push(small, nums[i])
    }
    for j := 0; j < k / 2; j ++ {
            x := heap.Pop(small)
            heap.Push(big, x)
    }
    if k % 2 == 0 {
        ans = append(ans, float64(small.Top() + big.Top()) / float64(2))
    } else {
        ans = append(ans, float64(small.Top()))
    }

    // mp保存应该出small或者big,但是没有出的数及其次数
    mp := make(map[int]int)
	for i := k; i < len(nums); i++ {
        // balance保存small比big中多出来数的个数
        balance := 0
        // nums[i - k]离开,mp[nums[i - k]] ++
        // nums[i - k] 小于等于small.Top(),balance -- 否则balance ++
		left := i - k 
        mp[nums[left]] ++

        if !small.Empty() && nums[left] <= small.Top() {
            balance --
        } else {
            balance ++
        }

        // mp[nums[i]] ++
        // nums[i] 小于等于small.Top(),balance -- 否则balance ++
        if !small.Empty() && nums[i] <= small.Top() {
            balance ++
            heap.Push(small, nums[i])
        } else {
            balance --
            heap.Push(big, nums[i])
        }

        // nums[left]离开双堆,nums[i]进入双堆,balance为0,-2,或者2
        // 考虑双堆中待删除的数,保证删除的数离开双堆后,始终保持双堆平衡
        // 一进一出之后,small比big多两个数
        if balance == 2 {
            heap.Push(big, heap.Pop(small))
        // 一进一出之后,small比big少两个数
        } else if balance == -2 {
            heap.Push(small, heap.Pop(big))
        }
        for !small.Empty() && mp[small.Top()] > 0 {
            mp[small.Top()] --
            heap.Pop(small)
        }
        for !big.Empty() && mp[big.Top()] > 0 {
            mp[big.Top()] --
            heap.Pop(big)
        }

        // 计算中位数,k为奇数从small中取,否则从big中取
        if k % 2 == 0 {
            ans = append(ans, float64(small.Top() + big.Top()) / float64(2))
        } else {
            ans = append(ans, float64(small.Top()))
        }
	}
	return ans
}


type SmallHeap []int
func (h SmallHeap) Len() int           { return len(h) }
func (h SmallHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h SmallHeap) Swap(i, j int)      { 
    h[i], h[j] = h[j], h[i] 
}
func (h *SmallHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

// pop是干掉最后一个元素,因为第一个和最后一个元素交换了位置,然后再执行pop,sink
func (h *SmallHeap) Pop() interface{} {
    old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}
// top是第一个元素
func (h *SmallHeap) Top() int {
    c := *h
    return c[0]
}
func (h *SmallHeap) Empty() bool { 
    return len(*h) == 0
}

type BigHeap []int
func (h BigHeap) Len() int           { return len(h) }
func (h BigHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h BigHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *BigHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}
func (h *BigHeap) Pop() interface{} {
    old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}
func (h *BigHeap) Top() int {
    c := *h
    return c[0]
}

func (h *BigHeap) Empty() bool { 
    return len(*h) == 0
}