返回

使用Go实现堆排序

前言

最近准备寻找下一份工作,也是到了再刷刷算法题的时候了,正巧碰到了堆排序的问题,虽说在Java中又优先队列可以避免手写,但想来再复习一下堆排序也无伤大雅,也是想着用Go来实现一下这个算法。

堆排序

基本概念

堆排序是一种基于二叉堆数据结构的原地排序算法,时间复杂度 O(nlogn),额外空间复杂度 O(1)。下面解释一下几个概念:

  • 完全二叉树:除了最后一层,所有层都是满的,且最后一层节点从左向右排列
  • 大顶堆:父节点大于等于子节点
  • 小顶堆:父节点小于等于子节点

算法执行步骤

首先需要构建堆(假设我们要构建一个大顶堆),步骤如下:

  1. 从最后一个非叶子节点开始 (n / 2 - 1)
  2. 对每个节点执行“下沉”操作,比较父子节点,将最大值交换给父节点
  3. 不断重复这个过程,直到根节点

之后就是根据这个堆进行排序,步骤如下:

  1. 将堆顶元素(最大值)与堆尾交换
  2. 堆大小减1
  3. 对堆顶执行下沉操作恢复大顶堆性质
  4. 不断重复直到堆大小为1

这里使用网上很热门的一个gif图来演示

堆排序gif图
堆排序gif图

优缺点

堆排序算法在最坏情况下,复杂度也是O(nlogn),且为原地排序,空间效率高。适合外部排序(数据在磁盘中)。缺点是属于不稳定排序。

算法实现

堆排序实现

整体实现的代码如下:

package main

import "fmt"

// 堆排序
func heapSort(arr []int) {
	n := len(arr) // 数组长度

	// 构建最大堆 从最后一个非叶子节点开始
	for i := n/2 - 1; i >= 0; i-- {
		heapify(arr, n, i)
	}

	// 排序
	for i := n - 1; i > 0; i-- {
		// 将堆顶的最大值移至数组末尾
		arr[i], arr[0] = arr[0], arr[i]
		// 从堆顶开始向下调整堆
		heapify(arr, i, 0)
	}
}

// 以索引i为根的子树调整为最大堆
func heapify(arr []int, n, i int) {
	largest := i
	left := i*2 + 1
	right := i*2 + 2

	// 判断当前三个节点的最大值
	if left < n && arr[left] > arr[largest] { // 如果要实现小顶堆 只需要调整第二个条件的比较符即可
		largest = left
	}
	if right < n && arr[right] > arr[largest] {
		largest = right
	}

	// 如果当前节点不是最大节点 交换 并向下调整堆
	if largest != i {
		arr[largest], arr[i] = arr[i], arr[largest]
		heapify(arr, n, largest)
	}
}

func main() {
	arr := []int{4, 12, 53, 11, 1, 65, 32}
	heapSort(arr)
	fmt.Println(arr)
}
====================
// 输出
[1 4 11 12 32 53 65]

首先需要从最后一个非叶子节点开始构建大顶堆。当构建完毕后,不断交换数组第一个和最后一个元素并恢复堆性质实现排序。当然,如果只是这样的实现就显得没什么意思,所以考虑使用泛型来实现这个堆排序。

泛型堆排序实现

package main

import "fmt"

// 使用 T 作为泛型变量 同时传递一个比较函数作为参数
func heapSort[T any](arr []T, less func(a, b T) bool) {
	n := len(arr)

	for i := n/2 - 1; i >= 0; i-- {
		heapify(arr, n, i, less)
	}

	for i := n - 1; i > 0; i-- {
		arr[i], arr[0] = arr[0], arr[i]
		heapify(arr, i, 0, less)
	}
}

func heapify[T any](arr []T, n, i int, less func(a, b T) bool) {
	largest := i
	left := 2*i + 1
	right := 2*i + 2

	if left < n && less(arr[largest], arr[left]) {
		largest = left
	}
	if right < n && less(arr[largest], arr[right]) {
		largest = right
	}
	if largest != i {
		arr[largest], arr[i] = arr[i], arr[largest]
		heapify(arr, n, largest, less)
	}
}
func main() {
	arr1 := []int{4, 12, 53, 11, 1, 65, 32}
	heapSort(arr1, func(a, b int) bool {
		return a < b
	})
	fmt.Println(arr1)

	arr2 := []float64{3.14, 1.41, 2.71, 1.62}
	heapSort(arr2, func(a, b float64) bool {
		return a < b
	})
	fmt.Println(arr2)
}
=================
// 输出
[1 4 11 12 32 53 65]
[1.41 1.62 2.71 3.14]

在这个实现中,整体实现流程与第一个版本相差无几,只不过是在传入的类型上添加了泛型T,并使用一个函数参数作为泛型变量比较的依据,在最后使用时传入匿名函数进行调用即可。根据上面的写法,同样也可以实现结构体的排序。

结构体堆排序

package main

import "fmt"

type Person struct {
	Name   string
	Age    int
	Height int
}

// 自定义结构体排序
func (p Person) Less(comPerson Person) bool {
	if p.Age != comPerson.Age {
		return p.Age < comPerson.Age
	}
	return p.Height < comPerson.Height
}

func heapSort(arr []Person) {
	n := len(arr)

	for i := n/2 - 1; i >= 0; i-- {
		heapify(arr, n, i)
	}

	for i := n - 1; i > 0; i-- {
		arr[0], arr[i] = arr[i], arr[0]
		heapify(arr, i, 0)
	}
}

func heapify(arr []Person, n, i int) {
	largest := i
	left := i*2 + 1
	right := i*2 + 2

	if left < n && arr[largest].Less(arr[left]) {
		largest = left
	}
	if right < n && arr[largest].Less(arr[right]) {
		largest = right
	}
	if largest != i {
		arr[i], arr[largest] = arr[largest], arr[i]
		heapify(arr, n, largest)
	}
}

func main() {
	arr := []Person{
		{"张三", 25, 175},
		{"李四", 30, 180},
		{"王五", 25, 170},
		{"赵六", 28, 185},
		{"钱七", 30, 175},
	}
	heapSort(arr)
	for _, person := range arr {
		fmt.Println(person)
	}
}
====================
// 输出
{王五 25 170}
{张三 25 175}
{赵六 28 185}
{钱七 30 175}
{李四 30 180}

在这个排序的实现中,额外定义了一个比较函数用来定义结构体之间的排序方式。同时传递的变量类型为Person,其实依旧可以使用泛型作为传递的参数。

heap包实现

当然,作为一门算是比较有名气的语言,go其实也有自己的heap实现。但不得不说的是,个人认为这个实现属实是有点繁琐了,希望之后go能够在这方面稍微改善一下吧。

package main

import (
	"container/heap"
	"fmt"
)

// 定义heap的底层数组
type IntHeap []int

// 实现heap 定义的方法
func (h IntHeap) Len() int {
	return len(h)
}
func (h IntHeap) Less(i, j int) bool {
	return h[i] < h[j]
}
func (h IntHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}
func (h *IntHeap) Push(x any) {
	*h = append(*h, x.(int)) // 需要进行类型断言
}
func (h *IntHeap) Pop() any {
	n := len(*h)
	x := (*h)[n-1]
	*h = (*h)[:n-1]
	return x
}

func main() {
	arr := &IntHeap{4, 12, 53, 11, 1, 65, 32}
	heap.Init(arr)
	heap.Push(arr, 23)
	fmt.Println("堆顶元素", (*arr)[0])
	for arr.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(arr))
	}
}
===========================
// 输出
堆顶元素 1
1 4 11 12 23 32 53 65

以上就是使用go语言实现堆的方式,当然这是一个固定的流程,没什么可以描述的。至于为什么这么实现,那就要深究到heap实现的源码了。

heap源码

接口定义

可以在 Go Packages 网址中找到heap包对应的源码。进入heap.go中首先可以看到,对于heap接口的定义

type Interface interface {
	sort.Interface
	Push(x any) // add x as element Len()
	Pop() any   // remove and return element Len() - 1.
}

因此如果想要实现堆,就得实现接口中的这几个方法。这个接口通过组合的方式引入了 sort.Interface ,进一步深入sort包可以看到,sort包中也有类似的接口定义:

type Interface interface {
	// Len is the number of elements in the collection.
	Len() int

	// Less reports whether the element with index i
	// must sort before the element with index j.
	//
	// If both Less(i, j) and Less(j, i) are false,
	// then the elements at index i and j are considered equal.
	// Sort may place equal elements in any order in the final result,
	// while Stable preserves the original input order of equal elements.
	//
	// Less must describe a [Strict Weak Ordering]. For example:
	//  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
	//  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
	//
	// Note that floating-point comparison (the < operator on float32 or float64 values)
	// is not a strict weak ordering when not-a-number (NaN) values are involved.
	// See Float64Slice.Less for a correct implementation for floating-point values.
	//
	// [Strict Weak Ordering]: https://en.wikipedia.org/wiki/Weak_ordering#Strict_weak_orderings
	Less(i, j int) bool

	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

这也就是使用heap时要实现的五个方法的由来了。在这个接口定义中, Interface 也就是堆对象,在go语言中,接口是隐式实现的,只要实现了接口中定义的所有方法,那么就默认实现了这个接口。所以当我们实现了这个上述这两个接口的方法,那么就已经一个堆了。

Init方法

下面来看一下初始化方法

func Init(h Interface) {
	// heapify
	n := h.Len()
	for i := n/2 - 1; i >= 0; i-- {
		down(h, i, n)
	}
}

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 指向右子节点
			j = j2 // = 2*i + 2  // right child
		}
		if !h.Less(j, i) { // 如果子节点不比父节点小,说明堆属性已满足
			break
		}
		h.Swap(i, j)
		i = j
	}
	return i > i0
}

稍微观察一下可以发现,其实与我们自己实现的堆排序算法差别不大,不过是将交换,比较逻辑抽象为函数罢了。还有一点不一样,那就是函数的返回值是布尔类型,如果返回 true 那说明发生了元素下沉。

为什么会有返回值,那是因为heap包提供了一个 Fix 方法用来修复堆。当堆中的某个位置的元素发生变化,并且堆的容量保持不变时,可以使用这个方法来修复堆。

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
	}
}

Push方法

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

相对前面的几个方法来说, Push 方法就简单很多了,首先调用我们自己实现的 Push 逻辑(也就是在切片末尾添加元素),之后再尝试将元素上移来保证堆的性质。

Pop方法

最后来看看Pop方法

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

可以看到,Pop方法首先会将切片首尾元素互换,之后将当前根节点的元素下沉,最后调用自定义的 Pop 方法来弹出元素。

这里也就很好的解释了为什么之前调用以下代码打印的时候得到的是一样的结果了

fmt.Println("堆顶元素", (*arr)[0])
for arr.Len() > 0 {
    fmt.Printf("%d ", heap.Pop(arr))
}

这是因为不调用 Pop 方法时,整体切片就已经是一个有序的状态了,在调用 Pop 时,会将元素首尾交换,因此实际上打印的是切片的第一个元素而不是最后一个元素。

到这里也就是完整的了解了整个heap的实现,但还是希望之后的版本能够在实现heap时能够舍弃掉这些繁琐的模板代码。

你要相信流星划过会带给我们幸运,就像现实告诉你我要心存感激