【Go学习笔记】堆（container/heap）

源码

package heap

import "sort"

// The Interface type describes the requirements
// for a type using the routines in this package.
// Any type that implements it may be used as a
// min-heap with the following invariants (established after
// Init has been called or if the data is empty or sorted):
//
//	!h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
//
// Note that Push and Pop in this interface are for package heap's
// implementation to call. To add and remove things from the heap,
// use heap.Push and heap.Pop.
type Interface interface {
	sort.Interface
	Push(x interface{}) // add x as element Len()
	Pop() interface{}   // remove and return element Len() - 1.
}

// Init establishes the heap invariants required by the other routines in this package.
// Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// The complexity is O(n) where n = h.Len().
func Init(h Interface) {
	// heapify
	n := h.Len()
	for i := n/2 - 1; i >= 0; i-- {
		down(h, i, n)
	}
}

// Push pushes the element x onto the heap.
// The complexity is O(log n) where n = h.Len().
func Push(h Interface, x interface{}) {
	h.Push(x)
	up(h, h.Len()-1)
}

// Pop removes and returns the minimum element (according to Less) from the heap.
// The complexity is O(log n) where n = h.Len().
// Pop is equivalent to Remove(h, 0).
func Pop(h Interface) interface{} {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

// Remove removes and returns the element at index i from the heap.
// The complexity is O(log n) where n = h.Len().
func Remove(h Interface, i int) interface{} {
	n := h.Len() - 1
	if n != i {
		h.Swap(i, n)
		if !down(h, i, n) {
			up(h, i)
		}
	}
	return h.Pop()
}

// Fix re-establishes the heap ordering after the element at index i has changed its value.
// Changing the value of the element at index i and then calling Fix is equivalent to,
// but less expensive than, calling Remove(h, i) followed by a Push of the new value.
// The complexity is O(log n) where n = h.Len().
func Fix(h Interface, i int) {
	if !down(h, i, h.Len()) {
		up(h, i)
	}
}

func up(h Interface, j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		j = i
	}
}

func down(h Interface, i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
			j = j2 // = 2*i + 2  // right child
		}
		if !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		i = j
	}
	return i > i0
}

分析

heap.Interface接口需要实现：

• Push(x interface{})
• Pop() interface{}

由于嵌入了sort.Interface接口，需要另外实现其声明的方法集：

• Len() int
• Less(i, j int) bool
• Swap(i, j int)

另外，标准库中也包含了如下函数：

• Push(h Interface, x interface{})
• Pop(h Interface) interface{}

分析函数的实现：

• Push()函数调用接口中的Push()方法向堆中添加元素，再调用up()函数自底向上调整堆
• Pop()函数交换堆顶与最后一个元素，调用down()函数自顶向下调整堆，最后调用接口中的Pop()方法弹出堆顶元素

func Push(h Interface, x interface{}) {
	h.Push(x)
	up(h, h.Len()-1)
}

func Pop(h Interface) interface{} {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

通过分析可知：

1. 在实现heap.Interface接口时，对于Push()和Pop()两个方法的实现，只需要包括元素的加入和弹出，不包括堆的调整
2. Pop()函数中，先调用了Swap()方法交换堆顶与末尾元素，因此在实现Pop()方法时，需要注意返回的应是末尾元素（原来的堆顶），在实现Push()方法时也是向末尾加入新元素
3. 在使用标准库时，应该调用库中的heap.Push()和heap.Pop()函数来进行堆操作，而非自己实现的方法

再看up()和down()两个主要堆调整函数：

func up(h Interface, j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		j = i
	}
}

func down(h Interface, i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
			j = j2 // = 2*i + 2  // right child
		}
		if !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		i = j
	}
	return i > i0
}

分析可知：库中实现的默认是一个小顶堆（调用Less()方法），因此若我们需要实现一个大顶堆，只需要在实现Less()方法时进行调整即可

实例

LC-2208

//构造大顶堆

type MaxHeap []float64

func (h MaxHeap) Len() int {
	return len(h)
}

func (h MaxHeap) Less(i, j int) bool {
	return h[i] > h[j] //返回较大者
}

func (h MaxHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *MaxHeap) Push(x interface{}) {
	*h = append(*h, x.(float64)) //接口类型转换
}

func (h *MaxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func halveArray(nums []int) int {
	n := len(nums)
	h := &MaxHeap{}
	heap.Init(h)
	var sum, cnt int
	var minus float64
	for i := 0; i < n; i++ {
		//统计初始数组之和
		sum += nums[i]
		//加入堆
		heap.Push(h, float64(nums[i]))
	}
	//将堆顶元素减半后重新插入
	for float64(sum) > 2*minus {
		top := heap.Pop(h).(float64) //接口类型转换
		top /= 2.0
		minus += top
		heap.Push(h, top)
		cnt++
	}
	return cnt
}

剑指 Offer 49

type hp struct{
    sort.IntSlice
}

func(h *hp)Push(v interface{}){
    h.IntSlice=append(h.IntSlice,v.(int))
}

func(h* hp)Pop()interface{}{
    n:=len(h.IntSlice)-1
    top:=h.IntSlice[n]
    h.IntSlice=h.IntSlice[:n]
    return top
}
var nums=[]int{2,3,5}
func nthUglyNumber(n int) int {
    h:=&hp{sort.IntSlice{1}}
    vis:=map[int]struct{}{1:{}}
    top:=1
    for i:=1;i<=n;i++{
        top=heap.Pop(h).(int)
        for j:=0;j<len(nums);j++{
            mul:=nums[j]*top
            if _,ok:=vis[mul];!ok{
                heap.Push(h,nums[j]*top)
                vis[mul]=struct{}{}
            }
        }
    }
    return top
}