【golang】基于go语言的堆结构模板

模板

package heap

type Heap struct {
	data    []int
	size    int
	compare func(int, int) int
}

//构造函数
func Construct(data []int, compare func(int, int) int) *Heap {
	heap := &Heap{data: data, size: len(data), compare: compare}
	heap.heapify()
	return heap
}

//获取父节点
func parent(node int) int {
	return (node - 1) / 2
}

//获取左子节点
func left(node int) int {
	return node*2 + 1
}

//获取右子节点
func right(node int) int {
	return node*2 + 2
}

//下沉
func (this *Heap) sink(node int) {
	//空节点
	if node >= this.size {
		return
	}
	//叶子节点
	if left(node) >= this.size {
		return
	}
	//右节点为空，只需判断父节点和左节点需不需要交换即可
	if right(node) >= this.size {
		if this.compare(this.data[node], this.data[left(node)]) > 0 {
			this.data[node], this.data[left(node)] = this.data[left(node)], this.data[node]
		}
		return
	}
	//说明当前节点为根的堆符合堆结构，无需调整
	if this.compare(this.data[node], this.data[left(node)]) <= 0 && this.compare(this.data[node], this.data[right(node)]) <= 0 {
		return
	}
	if this.compare(this.data[left(node)], this.data[right(node)]) <= 0 {
		this.data[node], this.data[left(node)] = this.data[left(node)], this.data[node]
		this.sink(left(node))
	} else {
		this.data[node], this.data[right(node)] = this.data[right(node)], this.data[node]
		this.sink(right(node))
	}
}

//上浮
func (this *Heap) floor(node int) {
	if node == 0 {
		return
	}
	if this.compare(this.data[node], this.data[parent(node)]) < 0 {
		this.data[node], this.data[parent(node)] = this.data[parent(node)], this.data[node]
		this.floor(parent(node))
	}
}

//堆调整，用于初始化
func (this *Heap) heapify() {
	now := parent(this.size - 1)
	for ; now >= 0; now-- {
		this.sink(now)
	}
}

//弹出堆顶
func (this *Heap) Pop() int {
	if this.size == 0 {
		panic("heap size is zero")
	}
	ans := this.data[0]
	this.data[0], this.data[this.size-1] = this.data[this.size-1], this.data[0]
	this.size--
	this.sink(0)
	return ans
}

//添加元素
func (this *Heap) Push(value int) {
	if this.size == len(this.data) {
		this.data = append(this.data, value)
	} else {
		this.data[this.size] = value
	}
	this.size++
	this.floor(this.size - 1)
}

//批量添加元素
func (this *Heap) PushMany(value ...int) {
	for _, v := range value {
		this.Push(v)
	}
}

测试

func main() {
	h1 := heap.Construct([]int{6, 1, 4, 9, -5, 8, 2, 7, 6, 11, 4}, func(a int, b int) int {
		return a - b
	})
	fmt.Println(h1)
	h2 := heap.Construct([]int{}, func(a int, b int) int {
		return a - b
	})
	h2.PushMany(6, 1, 4, 9, -5, 8, 2, 7, 6, 11, 4)
	fmt.Println(h2)
}

输出
4
5