TopK和堆排序

 

在处理海量数据或者数据流里经常会碰到提取前K个最大值，前K个最小值的问题；面试的时候也特别喜欢问此种问题。

本文先介绍了大顶堆和小顶堆，然后给出利用大顶堆和小顶堆算法解决TopK问题，最后给出了堆排序。

所有代码位于TopK

大顶堆和小顶堆

大顶堆和小顶堆都是完全二叉树，由于可以使用数组nums表示完全二叉树，所以存储结构用数组很方便。

父节点索引为i，则左孩子节点索引为2*i+1，右孩子节点索引为2*i+2

根节点索引为0，索引为1的节点为根节点的左孩子，索引为2的节点为根节点的右孩子。

即值为9节点的左孩子值为8，右孩子值为7。

大顶堆(Max-Heap)

每个结点的值都大于或等于其左右孩子结点的值，即nums[i]>=nums[2*i+1] && nums[i]>=nums[2*i+2]。

可以看到堆顶即根节点的值最大；

大顶堆可以用来找前K个最小值。

假设堆里面应该有K个值，其中根节点值最大，此时新来一个值。

  1. 如果比根节点大，则说明比堆里面其他节点都大，肯定不在前K个最小值里面；
  2. 如果比根节点小，则弹出最大的值，放入该值，即用该值替换根节点，然后重新调整形成大顶堆，用于下一次值调整。

同时大顶堆可用于升序堆排序，将大顶堆的根节点值与最后一个节点值交换，剩余的元素在构建大顶堆，依次构建形成升序。具体可以看后面的堆排序。

小顶堆(Min-Heap)

每个结点的值都大于或等于其左右孩子结点的值，即nums[i]<=nums[2*i+1] && nums[i]<=nums[2*i+2]。

同大顶堆，小顶堆用来找前K个最大值，用于堆排序的降序。

堆化，以大顶堆为例

假设树为

1
2
3
4
5
6
7
8
9
10
int array[] = {4, 9, 5, 7, 6, 8};
vector<int> nums(array, array + sizeof(array)/sizeof(int));
//对nums进行堆化，形成大顶堆
heapify(nums);
//输出堆化之后的
for (size_t i = 0; i < nums.size(); i++)
{
    cout << nums[i] << "\t";
}
cout << endl;

输出结果

9 7 8 4 6 5

对应的完全二叉树为

大顶堆的每一棵子树都是大顶堆，适用于递归算法；算法heapify的核心思想，从底向上，使得每棵树包括子树都变成大顶堆；类似于冒泡排序，从父节点和左右孩子节点选择最大的值作为父节点的值，由于是从底向上，能够保证根节点是最大值；

1
2
3
4
5
6
7
8
9
10
void heapify(vector<int> &nums)
{
    int array_size = nums.size();
    //(array_size-1)/2第一个非叶子节点
    //自底向上遍历所有的非叶子节点
    for (int i = (array_size - 1) / 2; i >= 0; i--)
    {
        build_heap(nums, i);
    }
}

算法build_heap会从父节点i开始构建大顶堆，由于是自底向上构建的，父节点i的左右孩子节点都已经是大顶堆了。如果父节点的值比左右孩子节点的值都大，则不用调整该树已经是大顶堆了；否则，交换父节点和左右孩子值最大的节点，由于有可能使孩子变成的非大顶堆，就需要在调用build_heap调整一下该孩子节点再次成为大顶堆。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
void build_heap(vector<int> &nums, int parent)
{
    int left = 2 * parent + 1;
    int right = 2 * parent + 2;
    int largest = parent;
    int array_size = nums.size();
    //max-heap
    if (left < array_size && nums[left] > nums[largest])
    {
        largest = left;
    }
    if (right < array_size && nums[right] > nums[largest])
    {
        largest = right;
    }
    if (largest != parent)
    {
        std::swap(nums[largest], nums[parent]);
        //swap之后孩子节点可能不是大顶堆了，需要调整一下
        build_heap(nums, largest);
    }
}

对树[9, 4, 5, 7, 6, 8]进行调整，非叶子节点的索引为[2,1,0]

  1. 调整非叶子节点2，的节点2值为5，只有左孩子节点值为8；交换5，8，由于左孩子节点无孩子节点，不需要继续调整。

  2. 调整非叶子节点1，由于9>7&&9>6，所以该树不用调整

  3. 调整节点0，由于左孩子值最大，交换节点0，节点1的值

    由于节点1不再是满足大顶堆的性质，需要调整成为大顶堆

    调整之后的大顶堆为[9,7,8,4,6,5]

整体代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#include <vector>
#include <iostream>
using namespace std;

void build_heap(vector<int> &nums, int parent)
{
    int left = 2 * parent + 1;
    int right = 2 * parent + 2;
    int largest = parent;
    int array_size = nums.size();
    //max-heap
    if (left < array_size && nums[left] > nums[largest])
    {
        largest = left;
    }
    if (right < array_size && nums[right] > nums[largest])
    {
        largest = right;
    }
    if (largest != parent)
    {
        std::swap(nums[largest], nums[parent]);
        build_heap(nums, largest);
    }
}

void heapify(vector<int> &nums)
{
    int array_size = nums.size();
    for (int i = (array_size - 1) / 2; i >= 0; i--)
    {
        build_heap(nums, i);
    }
}

int main()
{
    int array[] = {9, 4, 5, 7, 6, 8};
    vector<int> nums(array, array + sizeof(array)/sizeof(int));
    heapify(nums);
    for (size_t i = 0; i < nums.size(); i++)
    {
        cout << nums[i] << "\t";
    }
    cout << endl;
    return 0;
}

TopK

有了上面的heapify构建TopK就容易多了。

以查找前K个最小值为例，使用大顶堆。固定堆大小为K，遍历所有元素。

如果堆大小小于K，添加元素到堆中，然后堆化，其中根节点是最大值；

如果堆大小为K，比较根节点与元素的大小。

  1. 如果比根节点大，则比堆中其他K-1个元素都大，所以不会出现前K个最小值中；
  2. 如果比跟节点小，用该元素替换堆中最大值，即根节点，然后调整根节点为大顶堆(孩子节点已经是大顶堆了)。

具体代码如下，

  1. 指定compare函数可以支持大顶堆或者小顶堆
  2. push函数先TopK中添加元素
  3. get函数获取前K个元素；
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
#ifndef TOPK_H
#define TOPK_H

#include <vector>

template<typename T>

class TopK {
public:

	TopK(int k, bool (*compare_func)(const T &t1, const T &t2)) {
		k_ = k;
		compare_ = compare_func;
		array_.reserve(k_);
	}

	bool push(const T &value);

	void heapify(int array_size);

	std::vector<T> get() {
		return array_;
	}

	void build_heap(int array_size, int root);

private:
	int k_;
	std::vector<T> array_;
	bool (*compare_)(const T &t1, const T &t2);
};

template<typename T>
void TopK<T>::heapify(int array_size) {
	for (int i = (array_size - 1) / 2; i >= 0; i--) {
		build_heap(array_size, i);
	}
}

template<typename T>
bool TopK<T>::push(const T &value) {
	if (array_.size() < k_) {
		//array has less than k_ elements
		array_.push_back(value);
		heapify(array_.size());
		return true;
	}

	//array has k_ elements
	//user define compare func
	if (compare_ == NULL)
		return false;
	//if heap is min-heap, compare is greater func
	//if heap is max-heap, compare is less func
	if (compare_(array_[0], value)) {
		//the value need to insert heap
		array_[0] = value;
        //0 for root
		build_heap(k_,0);
	}
	return true;
}

template<typename T>
void TopK<T>::build_heap(int array_size, int root) {
	int left = 2 * root + 1;
	int right = 2 * root + 2;
	int largest = root;
	if (left < array_size && compare_(array_[left], array_[largest])) {
		largest = left;
	}
	if (right < array_size && compare_(array_[right], array_[largest])) {
		largest = right;
	}
	if (largest != root) {
		std::swap(array_[largest], array_[root]);
		build_heap(array_size, largest);
	}
}

#endif //TOPK_H

测试代码，使用小顶堆找出[4, 5, 1, 9, 8, 0, 3, 2]中Top3最大值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#include <iostream>
#include "TopK.h"

//max-heap
bool max_heap_compare(const int &t1, const int &t2)
{
    return t1 > t2;
}
//min-heap
bool min_heap_compare(const int &t1, const int &t2)
{
    return t1 < t2;
}

int main()
{
    int array[] = {4, 5, 1, 9, 8, 0, 3, 2};
    int array_len = sizeof(array) / sizeof(array[0]);
    TopK<int> top_k(3, min_heap_compare);
    for (int i = 0; i < array_len; i++)
    {
        top_k.push(array[i]);
        std::vector<int> result = top_k.get();
        for (int i = 0; i < result.size(); i++)
        {
            std::cout << result[i] << " ";
        }
        std::cout << std::endl;
    }

    
    return 0;
}

程序输出

4
4 5
1 5 4
4 5 9
5 8 9
5 8 9
5 8 9
5 8 9

堆排序

堆排序是一种选择排序，它的最坏，最好，平均时间复杂度均为O(nlogn)，它也是不稳定排序

以升序为例，构建大顶堆。要排序的数组array大小为n

  1. 利用build_heap算法，对array[0,…,n-1]进行堆化，需要自底向上调整所有非叶子节点，使整棵树变成大顶堆，其中根节点为数组最大值，交换array[0]和array[n-1]
  2. 对array[o,…,n-2]进行堆化，这里只需要对堆化根节点，根节点为最大值，交换array[0]和array[n-2]
  3. 依次循环步骤2之至堆中只有一个元素，此时数据已经排好序
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#include <iostream>

void build_heap(int array[], int array_len, int root)
{
    int left = root * 2 + 1;
    int right = root * 2 + 2;
    int largest = root;
    if (left < array_len && array[left] > array[largest])
    {
        largest = left;
    }

    if (right < array_len && array[right] > array[largest])
    {
        largest = right;
    }

    if (largest != root)
    {
        std::swap(array[largest], array[root]);
        build_heap(array, array_len, largest);
    }
}

void heap_sort(int array[], int array_len)
{
    //first build heap make sure the tree is a heap
    for (int i = (array_len - 1) / 2; i >= 0; i--)
    {
        build_heap(array, array_len, i);
    }

    for (int i = array_len - 1; i >= 0; i--)
    {
        //put the max value to the last index
        std::swap(array[0], array[i]);
        build_heap(array, i, 0);
    }
}

void print_array(int array[], int array_len)
{
    for (int i = 0; i < array_len; i++)
    {
        std::cout << array[i] << " ";
    }
    std::cout << std::endl;
}

int main()
{
    int array[] = {6, 9, 8, 1, 4, 2, 3, 5, 7};
    int array_len = sizeof(array) / sizeof(array[0]);
    heap_sort(array, array_len);
    print_array(array, array_len);
    return 0;
}

程序输出结果

1 2 3 4 5 6 7 8 9

 上一頁

C/C++ makefile