【堆排｜快排】Top-k问题

引言

• 决解top-k问题一般有：堆和快排两种思路

  两种算法的优劣如下：

  • 快速排序算法的时间、空间复杂度都优于堆排序算法
  • 但是快速排序算法需要修改原数组，如果原来的数组不能修改的话，还需要拷贝一份数组，造成空间复杂度升高
  • 快速排序算法需要保存所有数据后才能选出top-k。如果把数据看成输入流的话，使用堆方法则是来一个处理一个，不需要保存数据，只需要保存top-k个元素的堆。所以当数据量非常大的时候（大到内存放不下的时候），使用快速排序算法就无法解决目标问题。
  • 堆排序算法适合处理数据量大的top-k问题

两种算法原理

大顶堆/小顶堆

本文以实现大顶堆的原理为例子

• 大顶堆的定义：
  根节点必须大于左右子树的根节点，左右子树也必须满足大顶堆的定义
• 小顶堆的定义：
  根节点必须小于左右子树的根节点，左右子树也必须满足小顶堆的定义
• 大顶堆和小顶堆的应用场景：
  求升序用大顶堆，求降序用小顶堆

构造大顶堆的实现思路：

1. 先找到初始二叉树的最后一个非叶子节点的位置，如果当前位置节点小于左孩子节点，则和左孩子节点的位置进行交换
2. 在第一步的基础上判断当前位置节点是否小于右孩子节点，如果小于则和后孩子节点的位置进行交换
3. 选择上一个非叶子节点重复步骤1，直到无前一个节点

构造流程：

建堆的时间复杂度是：O(n)

1. 最后一个非叶子节点为8，根据步骤1和步骤2将 8 和 13、16的位置进行相应的交换

2. 如此往复，最后成功构造大顶堆
   

大顶堆的排序流程

实现思路：
不断把堆顶元素和堆的最后一个数据进行位置交换，交换一次位置后就把已经排好序前面的n-1元素再次执行堆构造算法重新生成大顶堆。

Golang利用大顶堆算法实现top-k算法

以下是利用大顶堆算法实现top-k小的算法

package datastruct

import (
	"container/heap"
)

/*
golang 自带的heap实现自定义大顶堆/小顶堆
*/

type IntHeap []int

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

// Less 决定是大顶堆还是小顶堆，此处实现的是大顶堆
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] }

// 实现heap 的Swap接口
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

// 实现heap 的Pop接口
func (h *IntHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// 实现heap 的 Push接口
func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func getLeastNumbers(arr []int, k int) []int {
	h := &IntHeap{}
	heap.Init(h)
	for _, val := range arr {
    // 将新的数据加大顶堆中，并将堆进行排序
		heap.Push(h, val)
		if h.Len() > k {
      // 如果堆中的元素个数大于k，则把最后一个节点踢出堆
      //（由于大顶堆的排序，最后一个节点肯定是所有节点中最大的节点）
			heap.Pop(h)
		}
	}
	return []int(*h)
}

快速选择（快排变形）

在快速排序中用一步很重要的操作：partition（划分）操作，从数组中随机选取一个枢纽元素v ，然后原地移动数组中的元素，使得比v小大元素在v的左边，比v大的元素在v的右边

这个partition操作是原地进行的，需要O(n)的时间，接下来快排排序算法会递归左右两侧的数组。而快速选择算法相当于一个“不完全”的快速排序，因为我们只需要知道最小的k个数是哪些，并不需要知道它们的顺序。

假设经过一次partition操作，枢纽元素位于下表m，也就是说左侧的数组有m个元素，是原数组中最小的m个数那么：

• 如果k=m,我们就找到了最小的k个数，就是左侧的数组
• 如果k<m，则最小的k个数一定都在左侧数组中，我们只需对左侧数组递归地partition即可
• 如果k>m，则左侧数组中的m个数都属于最小的k个数，我们还需要在右侧数组中寻找最小k-m个数，对右侧数组递归地partition即可。

Golang利用大顶堆算法实现top-k算法

package datastruct

func QuickSortTopK(arr []int, k int) []int {
	if k == 0 {
		return []int{}
	}
	if k == len(arr) {
		return arr
	}
	partitionArray(arr, 0, len(arr)-1, k)
	return arr[:k]
}

func partitionArray(arr []int, low, high, k int) {
	// 第一次partition
	mid := partition(arr, low, high)
	if mid == k {
		return
	} else if mid > k {
		// 如果k<mid 证明top-k在mid中的左侧中
		partitionArray(arr, low, mid-1, k)
	} else if k > mid {
		// 如果k>mid 证明top-k有一部分在mid的右侧，则继续寻找k-m部分
		partitionArray(arr, mid+1, high, k)
	}

}
func partition(arr []int, low, high int) int {
	l, r := low, high+1
	mid := arr[low]
	for {
		// 所有小于mid的值被放在左侧
		l++
		for arr[l] < mid {
			l++
			if l == high {
				break
			}
		}
		// 所有大于mid的树被放在右侧
		r--
		for arr[r] > mid {
			r--
			if r == low {
				break
			}
		}
		// 越界时跳出循环
		if l >= r {
			break
		}
		// 大于mid的数放右侧，小于mid的数放左侧
		arr[l], arr[r] = arr[r], arr[l]
	}
	//因为最后r一定会指向小于mid孩子等于mid的值，所以开始的节点要和r最后指向的节点交换位置
	arr[low], arr[r] = arr[r], arr[low]
	// 返回分隔位置
	return r
}