最大堆的应用

在刷题时遇到一个适合使用最大堆来解决的问题，刚好在C++标准库的算法库中有堆操作的相关算法，可以直接应用在容器上。回忆一下最大堆的原理，简单实现一下堆操作。

215. Kth Largest Element in an Array

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

Given an integer array nums and an integer k, return the kth largest element in the array.
Note that it is the kth largest element in the sorted order, not the kth distinct element.
Can you solve it without sorting?

Example 1:
Input: nums = [3,2,1,5,6,4], k = 2
Output: 5

Example 2:
Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

思路
定义一个长度为k的vector，遍历nums里的每一个数，如果比vector中最小的数大，则替换掉最小的数。最后vector里最小的数就是nums里第k大的数。
这个长度为k的vector，需要快速的找到最小的数，则标准库里的堆操作可以用。

解题方法
熟悉标准库的算法库里的堆操作。cppreference 页面。

1
2
3
4
5
6

make_heap(first, last, comp);
把[first, end)范围内的数转换成最大堆，转换后*first就是最大值。如果comp传入std::greater()，则是转换成最小堆。
pop_heap(first, last, comp);
弹出最小值，实质是把最大值移到范围的最后一位。
push_heap(first, last, comp);
插入一个数，实质是把范围的最后一位数，插入到前面，维持一个最大堆。

代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        vector<int> knums(k, std::numeric_limits<int>::min());
        for (int n: nums) {
            if (n > knums[0]) {
                pop_heap(knums.begin(), knums.end(), std::greater());
                knums[k-1] = n;
                push_heap(knums.begin(), knums.end(), std::greater());
            }
        }
        return knums[0];
    }
};

最大堆实现

最大堆的定义是：

  • 是完全二叉树
  • 父节点比子节点大

试着实现标准库里的堆操作算法。把场景简化为一个int的vector，用来表示完全二叉树，然后在vector上实现堆操作。对于完全二叉树的数组表示，设完全二叉树的一元素编号为i，i>=1, i<=n，有如下性质：

  • i=1，元素是根节点。否则，其父节点编号是[i/2]。
  • 2i>n，则该元素无左孩子，否则，左孩子编号是2i。
  • 2i+1>n，则该元素无右孩子，否则，右孩子编号是2i+1。

push_heap

首先是比较简单的插入，只需要把新的数据作为完全二叉树的最后一个子节点插入，然后再一路向上与其父节点比较，若比父节点大，就交换父子节点即可。

 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

namespace MyHeapOp {
void swap(std::vector<int> &a, int i, int j) {
    int tmp = a[i];
    a[i] = a[j];
    a[j] = tmp;
}

void push_heap(std::vector<int> &a) {
    int n = a.size();
	while (n > 1) {
        // a[n-1] is child，a[n/2-1] is the parent
        if (a[n-1] > a[n/2-1]) {
            // swap child and parent
            swap(a, n-1, n/2-1)
        }
        n = n/2;
    }
}
}

void print_vec(std::vector<int> &a) {
    for (auto x : a) {
        std::cout << x << " ";
    }
    std::cout << "\n";
}

void test_push_heap() {
    std::vector<int> a{4, 6, 7, 2, 5, 3, 9};
    int n = a.size();
    std::cout << "push_heap\n";
    print_vec(a);

    std::cout << "MyHeapOp::push_heap\n";
    std::vector<int> b;
    b.reserve(n);
    for (auto x : a) {
        b.push_back(x);
        MyHeapOp::push_heap(b);
        print_vec(b);
    }

    std::cout << "std::push_heap\n";
    std::vector<int> c;
    c.reserve(n);
    for (auto x : a) {
        c.push_back(x);
        std::push_heap(c.begin(), c.end());
        print_vec(c);
    }
}

测试输出，与标准库的push_heap结果一样，没有问题。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

push_heap
4 6 7 2 5 3 9 
MyHeapOp::push_heap
4 
6 4 
7 4 6 
7 4 6 2 
7 5 6 2 4 
7 5 6 2 4 3 
9 5 7 2 4 3 6 
std::push_heap
4 
6 4 
7 4 6 
7 4 6 2 
7 5 6 2 4 
7 5 6 2 4 3 
9 5 7 2 4 3 6

pop_heap

最大堆的删除操作pop_heap是把最大元素放在vector的最后位置，然后前n-1个元素仍然是最大堆。其思路是把原来最后一个元素放在根节点，再从根节点开始，比较其两个左右子节点，选择较大者与父节点交换，继续下去，直到一个叶节点。

 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

namespace MyHeapOp
{
#define left_child(i) (2*(i)+1)
#define right_child(i) (2*(i)+2)
void swap(std::vector<int> &a, int i, int j) {
    int tmp = a[i];
    a[i] = a[j];
    a[j] = tmp;
}
void pop_heap(std::vector<int> &a) {
    int n = a.size();
    if (n == 0) return;
    int front = a[0];
    a[0] = a[n-1];
    int i = 0;
    while (left_child(i) < n) {
        int swap_idx;
        if (right_child(i) < n and a[right_child(i)] > a[left_child(i)]) {
            swap_idx = right_child(i);
        } else {
            swap_idx = left_child(i);
        }
        if (a[swap_idx] > a[i]) swap(a, i, swap_idx);
        i = swap_idx;
    }
    a[n-1] = front;
}
}

void test_pop_heap() {
    std::vector<int> a{9, 5, 7, 2, 4, 3, 6};
    std::cout << "pop_heap\n";
    print_vec(a);
    std::cout << "MyHeapOp::pop_heap\n";
    while (not a.empty()) {
        MyHeapOp::pop_heap(a);
        print_vec(a);
        a.pop_back();
    }
    std::vector<int> b{9, 5, 7, 2, 4, 3, 6};
    std::cout << "std::pop_heap\n";
        while (not b.empty()) {
        std::pop_heap(b.begin(), b.end());
        print_vec(b);
        b.pop_back();
    }
}

测试打印为

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

pop_heap
9 5 7 2 4 3 6 
MyHeapOp::pop_heap
7 5 6 2 4 3 9 
6 5 3 2 4 7 
5 4 3 2 6 
4 2 3 5 
3 2 4 
2 3 
2 
std::pop_heap
7 5 6 2 4 3 9 
6 5 3 2 4 7 
5 4 3 2 6 
4 2 3 5 
3 2 4 
2 3 
2

与标准库一致，成功。

make_heap

make_heap一种思路是使用一个空的vector，把原vector插入进去，像push_heap的测试一样，但是这样复杂度为O(nlogn)。可以直接在原vector上进行调整，从最后一个具有子节点的元素开始检查每个子树是否是最大堆，将其调整为最大堆。

 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

namespace MyHeapOp {
int get_larger_child(std::vector<int> &a, int n, int parent) {
    if (right_child(parent) < n) {
        return a[left_child(parent)] > a[right_child(parent)] ? left_child(parent) : right_child(parent);
    }
    if (left_child(parent) < n) {
        return left_child(parent);
    }
    return -1;
}

void make_heap_child_tree(std::vector<int> &a, int n, int root) {
    int child = get_larger_child(a, n, root);
    while (child != -1) {
        if (a[root] < a[child]) swap(a, root, child);
        root = child;
        child = get_larger_child(a, n, root);
    }
}

void make_heap(std::vector<int> &a) {
    int n = a.size();
    if (n <= 1) return;
    int parent = n/2-1;
    while (parent >= 0) {
        make_heap_child_tree(a, n, parent);
        parent--;
    }
}
}

void test_make_heap() {
    std::vector<int> a{4, 6, 7, 2, 5, 3, 9};
    std::cout << "make_heap\n";
    print_vec(a);
    std::cout << "MyHeapOp::make_heap\n";
    MyHeapOp::make_heap(a);
    print_vec(a);
    std::vector<int> b{4, 6, 7, 2, 5, 3, 9};
    std::cout << "std::make_heap\n";
    std::make_heap(b.begin(), b.end());
    print_vec(b);
}

测试结果

1
2
3
4
5
6

make_heap
4 6 7 2 5 3 9 
MyHeapOp::make_heap
9 6 7 2 5 3 4 
std::make_heap
9 6 7 2 5 3 4

priority_queue

在STL中，优先级队列是基于vector的实现的最大堆。参考priority_queue

参考

数据结构、算法与应用
cppreference