矩阵中战斗力最弱的 K 行

通过这道题学习一下 C++ 中常用的库函数:upper_boundlower_boundrbeginrenddistancegreater<int>()move

upper_bound 上界,返回大于某个数的第一个位置。
lower_bound 下界,返回大于等于某个数的第一个位置。

这两个函数用于在有序数组中查找某个元素出现的第一个位置最后一个位置的下一个位置。

1
2
3
4
5
6
7
8
9
vector<int> a = {1, 2, 2, 2, 3, 4};
auto p = upper_bound(a.begin(), a.end(), 2);
auto q = lower_bound(a.begin(), a.end(), 2);
/*
运行后,p、q 两个迭代器的指向如下:
1, 2, 2, 2, 3, 4
q p
*/
int d = distance(q, p); // d == 3 distance(开始, 结束),如果写反了 distance(p, q) == -3

distance(p, q) 用于计算两个迭代器 p、q 之间的距离。相当于 p - q。

std::greater 返回一个比较函数,可以用于改变 sortpriority_queue 的行为:

1
2
if a > b return true
return false

std::move 用于复制一个容器内的元素到另一个容器。参考 https://www.geeksforgeeks.org/stdmove-in-c/

1
2
3
4
5
vector <int> vec1 {1, 2, 3, 4, 5};
vector <int> vec2 {7, 7, 7, 7, 7};
// 将 vec1 容器内的前四个数字复制到 vec2 第二个位置
move(vec1.begin(), vec1.begin() + 4, vec2.begin() + 1);
// vec2 {7, 1, 2, 3, 4}

rbeginrend 迭代器是倒着走的

1
2
3
4
5
6
7
8
vector<int> a = {1, 2, 2, 2, 3, 4};
auto p = a.rbegin();
while (p != a.rend()) {
cout << *p << endl;
p++;
}

// 逆序输出:4 3 2 2 2 1

本题思路,遍历矩阵每一行,由于 1 都在 0 的前面,所以可以通过二分搜索最后一个 1 的位置,也就是这些 1 的总和。将 1 的个数和当前行数作为 pair 放入优先队列,最后从队列取出 k 个元素就可以了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class Solution {
public:
vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
vector<pair<int, int>> h;
for (int i = 0; i < mat.size(); i++) {
auto p = upper_bound(mat[i].rbegin(), mat[i].rend(), 0);
int d = mat[i].rend() - p;
h.push_back({d, i});
}
priority_queue q(greater<pair<int, int>>(), move(h));
vector<int> ans;
for (int i = 0; i < k; i++) {
ans.push_back(q.top().second);
q.pop();
}
return ans;
}
};