【算法】堆和优先队列
2023-08-24 17:59:46

堆是一棵树,若父节点的值不小于子节点的值,就是大根堆,若父节点的值不大于子节点的值,就是小根堆。常用的堆为二叉大根堆。

二叉堆是一个完全二叉树。所谓完全二叉树就是,每一层的节点都”左对齐”,具体参考完全二叉树

堆有几种操作:

  1. 插入。向二叉堆中插入节点,只需在最后一层插入,之后根据堆的性质和父节点的值比较大小,调整位置即可,该操作最多比较的次数和层数几乎一致,故复杂度为$O(log_2n)$。

  2. 修改。同样复杂度为$O(log_2n)$。

  3. 删除根节点。该操作通常直接将层序从左向右遍历的最后一个节点赋给根节点的值,删除最后一个节点,接着根节点和子节点值比较,进行调整,复杂度同样为$O(log_2n)$。

优先队列

优先队列是排好序的队列。其本质就是一个堆。为什么?将堆拉直,第二层节点置于第一层右侧,第三层节点置于第二层右侧,就实现了一个线性的数据结构,由于删除常常删除根节点,而插入常常在末尾添加,这就是一个队列。

在C++里,优先队列是一个大根堆,由vector实现,严格来说,优先队列是vector的一个适配器。如果对性能有要求,可以自己写一个堆。

由于是使用vector实现的,那么讨论几种操作的复杂度:

  1. 插入。只需要从左到右将待插入的元素和队列里元素逐个比较即可,复杂度为$O(n)$,注意,这里并非将整个队列排序,而是对单个元素进行排序,所以复杂度不是$O(nlogn)$。
  2. 删除。复杂度为$O(n)$。

适用于什么问题

由于优先队列和堆的有序性,其常常用于“前k个…”问题。

比如Leetcode703. Kth Largest Element in a Stream:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class KthLargest {
public:
priority_queue<int, vector<int> ,greater<int>> pq;
int k;
KthLargest(int k, vector<int>& nums) {
this->k = k;
for(int i : nums) {
pq.push(i);
if(pq.size() > k) {
pq.pop();
}
}
}

int add(int val) {
pq.push(val);
if(pq.size() > k) {
pq.pop();
}
return pq.top();
}
};

Leetcode973. K Closest Points to Origin:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
class compare {
public:
bool operator()(vector<int> a, vector<int> b) {
return (a[0] * a[0] + a[1] * a[1]) < (b[0] * b[0] + b[1] * b[1]);
}
};
public:
vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
vector<vector<int>> result;
priority_queue<vector<int>, vector<vector<int>>, compare> pq;
for(vector<int> p : points) {
pq.push(p);
if(pq.size() > k) {
pq.pop();
}
}
while(!pq.empty()) {
result.push_back(pq.top());
pq.pop();
}
return result;
}
};

快速选择算法

该算法适用于求解上述问题,其本质是分治算法(divide and conquer)。

该算法的思想是:逐步缩小选择范围。

具体的解释在Top K 问题的最优解 - 快速选择算法(Quickselect)这里解释的已经很清楚了,这里只是用C++实现一下相关代码。

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
#include <iostream>

using namespace std;

int Partition(int* arr, int l, int r) {
int pivot = arr[r], i = l;
for(int j = l; j < r; ++j) {
if(arr[j] <= pivot) {
swap(arr[i], arr[j]);
++i;
}
}
swap(arr[i], arr[r]);
return i;
}

int KthSmallest(int* arr, int l, int r, int k) {
int part_index = Partition(arr, l, r);
if(k - 1 == part_index - l) {
return arr[part_index];
} else if(k - 1 < part_index - l) {
return KthSmallest(arr, l, part_index - 1, k);
} else {
return KthSmallest(arr, part_index+1, r, k - part_index - 1 + l);
}
return -1;
}

int main(int argc, char **argv) {
int arr[5] = {1, 2, 3, 4, 5};
cout << KthSmallest(arr, 0, 4, 3);
return 0;
}

用这个方法做Leetcode215. Kth Largest Element in an Array会导致TLE。

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
class Solution {
public:
int partition(vector<int>& nums, int l, int r) {
int pivot_index = rand() % (r - l + 1) + l, i = l;
swap(nums[pivot_index], nums[r]);
int pivot = nums[r];
for(int j = l; j < r; ++j) {
if(nums[j] <= pivot) {
swap(nums[i], nums[j]);
++i;
}
}
swap(nums[i], nums[r]);
return i;
}

int findKthLargest(vector<int>& nums, int k) {
int index = partition(nums, 0, nums.size() - 1);
while(index != nums.size() - k) {
if(index > nums.size() - k) {
index = partition(nums, 0, index - 1);
} else {
index = partition(nums, index + 1, nums.size() - 1);
}
}
return nums[index];
}
};

远不如调用std::priority_queue来得舒适。