题目描述
给你一个大小为 m * n 的矩阵 mat,矩阵由若干军人和平民组成,分别用 1 和 0 表示。
请你返回矩阵中战斗力最弱的 k 行的索引,按从最弱到最强排序。
如果第 i 行的军人数量少于第 j 行,或者两行军人数量相同但 i 小于 j,那么我们认为第 i 行的战斗力比第 j 行弱。
军人 总是 排在一行中的靠前位置,也就是说 1 总是出现在 0 之前。
示例1:
输入: mat =
[[1,1,0,0,0],
[1,1,1,1,0],
[1,0,0,0,0],
[1,1,0,0,0],
[1,1,1,1,1]],
k = 3
输出: [2,0,3]
解释: 每行中的军人数目:
行 0 -> 2
行 1 -> 4
行 2 -> 1
行 3 -> 2
行 4 -> 5
从最弱到最强对这些行排序后得到 [2,0,3,1,4]
示例2:
输入: mat =
[[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]],
k = 2
输出: [0,2]
解释:
每行中的军人数目:
行 0 -> 1
行 1 -> 4
行 2 -> 1
行 3 -> 1
从最弱到最强对这些行排序后得到 [0,2,3,1]
提示
- m == mat.length
- n == mat[i].length
- 2 <= n, m <= 100
- 1 <= k <= m
- matrix[i][j] 不是 0 就是 1
方法一:朴素做法
解题思路
统计出每行军人的数目,然后对其进行排序,输出前 k 行的索引 。
代码
class Solution {
typedef pair<int, int> PII;
public:
vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
vector<int> ans(k);
int n = mat.size();
vector<PII> cnt(n);
for(int i = 0; i < mat.size(); i++)
{
cnt[i] = {count(mat[i].begin(), mat[i].end(), 1), i};
}
sort(cnt.begin(), cnt.end(), [&](const auto& x, const auto& y){
return x.first == y.first ? x.second < y.second : x.first < y.first;
});
for(int i = 0; i < k; i++)
ans[i] = cnt[i].second;
return ans;
}
};
复杂度分析
- 时间复杂度:遍历矩阵的复杂度为 O ( m ∗ n ) O(m *n) O(m∗n);排序复杂度为 O ( m l o g m ) O(mlogm) O(mlogm);构造答案复杂度为 O ( k ) O(k) O(k)。整体复杂度为 O ( max ( m ∗ n , m log m ) ) O(\max(m * n, m\log{m})) O(max(m∗n,mlogm))。
- 空间复杂度: O ( m ) O(m) O(m) 空间用于存储所有的行战力; O ( l o g m ) O(logm) O(logm) 空间用于排序。整体复杂度为 O ( m + l o g m ) O(m + logm) O(m+logm)。
方法二:二分查找 + 堆
解题思路
军人总是排在一行中的靠前位置,也就是说 1 总是出现在 0 之前。因此,我们用二分查找来找到每一行中 1 出现的最后一个位置,下标 + 1 便是每行中 1 的个数,即每行的战斗力。
当我们得到每一行的战斗力后,我们可以将它们全部放入一个小根堆中,并不断地取出堆顶的元素 k 次,这样我们就得到了最弱的 k 行的索引。
如果我们依次将每一行的战斗力以及索引(因为如果战斗力相同,索引较小的行更弱,所以我们需要在小根堆中存放战斗力和索引的二元组)放入小根堆中,那么这样做的时间复杂度是
O
(
m
log
m
)
O(m \log m)
O(mlogm) 的。
一种更好的方法是使用这 m 个战斗力值直接初始化一个小根堆,时间复杂度为 O(m)。即
p
r
i
o
r
i
t
y
_
q
u
e
u
e
q
(
g
r
e
a
t
e
r
<
p
a
i
r
<
i
n
t
,
i
n
t
>
>
(
)
,
m
o
v
e
(
p
o
w
e
r
)
)
priority \_ queue\ q(greater<pair<int, int>>(), move(power))
priority_queue q(greater<pair<int,int>>(),move(power))
代码
class Solution {
public:
vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
int m = mat.size(), n = mat[0].size();
vector<pair<int, int>> power;
for(int i = 0; i < m; i++)
{
int l = 0, r = n - 1, mid;
while(l < r)
{
mid = (l + r + 1) / 2;
if(mat[i][mid] >= 1) l = mid;
else r = mid - 1;
}
power.emplace_back(l + 1, i);
}
// 数组直接初始化小根堆
priority_queue q(greater<pair<int, int>>(), move(power));
vector<int> ans;
for(int i = 0; i < k; i++)
{
ans.push_back(q.top().second);
q.pop();
}
return ans;
}
};
复杂度分析
- 时间复杂度:O(
m
l
o
g
n
+
k
l
o
g
m
mlogn+klogm
mlogn+klogm)。
需要 O ( m log n ) O(m \log n) O(mlogn) 的时间对每一行进行二分查找。
需要 O(m) 的时间建立小根堆。
需要 O ( k log m ) O(k \log m) O(klogm) 的时间从堆中取出 k 个最小的元素。 - 空间复杂度:O(m),即为堆所需要的空间。