leetcode 1337. The K Weakest Rows in a Matrix

Posted 2021-03-10 琴影

Given a m * n matrix mat of ones (representing soldiers) and zeros (representing civilians), return the indexes of the k weakest rows in the matrix ordered from the weakest to the strongest.

A row i is weaker than row j, if the number of soldiers in row i is less than the number of soldiers in row j, or they have the same number of soldiers but i is less than j. Soldiers are always stand in the frontier of a row, that is, always ones may appear first and then zeros.

 

Example 1:

Input: 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
Output: [2,0,3]
Explanation: 
The number of soldiers for each row is: 
row 0 -> 2 
row 1 -> 4 
row 2 -> 1 
row 3 -> 2 
row 4 -> 5 
Rows ordered from the weakest to the strongest are [2,0,3,1,4]

Example 2:

Input: mat = 
[[1,0,0,0],
 [1,1,1,1],
 [1,0,0,0],
 [1,0,0,0]], 
k = 2
Output: [0,2]
Explanation: 
The number of soldiers for each row is: 
row 0 -> 1 
row 1 -> 4 
row 2 -> 1 
row 3 -> 1 
Rows ordered from the weakest to the strongest are [0,2,3,1]

 

Constraints:

• m == mat.length
• n == mat[i].length
• 2 <= n, m <= 100
• 1 <= k <= m
• matrix[i][j] is either 0 or 1.

题目大意：

给你一个大小为 m * n 的方阵 mat，方阵由若干军人和平民组成，分别用 0 和 1 表示。

请你返回方阵中战斗力最弱的 k 行的索引，按从最弱到最强排序。

如果第 i 行的军人数量少于第 j 行，或者两行军人数量相同但 i 小于 j，那么我们认为第 i 行的战斗力比第 j 行弱。

军人 总是 排在一行中的靠前位置，也就是说 1 总是出现在 0 之前。

题目解析：

每一行的战斗力很容易求出来，但是要返回的是战斗力对应的行号，所以我们需要将每行的编号和行战斗力绑定起来，以便根据战斗力排序后很容易取出对应的行号。

绑定方式多种多样，二维数组、结构体或者pair都是可以的。

C++代码一：

 1 vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
 2         vector<int> ans(k);
 3         int m = mat.size();
 4         vector<array<int, 2>> v(m);
 5         for (int i = 0; i < m; ++i) {
 6             int c = 0;
 7             for (int j = 0; j < mat[i].size(); ++j) {
 8                 c += mat[i][j];
 9             }
10             v[i] = {c, i};
11         }
12         sort(v.begin(), v.end());
13         for (int i = 0; i < k; ++i) {
14             ans[i] = v[i][1];
15         }
16         return ans;
17     }

C++代码二：

 1 vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
 2         vector<int> ans(k);
 3         int r = mat.size(), c = (mat.size() == 0 ? 0 : mat[0].size()), index = 0;
 4         vector<bool> visited(r, false); //判断某一行号是否被选择
 5         for (int j = 0; j < c; ++j) {
 6             for (int i = 0; i < r; ++i) {
 7                 if (visited[i]) //如果第i行已经被选择，则不需要考虑这一行的军人数
 8                     continue;
 9                 if (index >= k) //提前找到了k行，则跳出两层循环
10                     goto a;
11                 if (index < k && mat[i][j] == 0) {
12                     ans[index++] = i;
13                     visited[i] = true;
14                 }
15             }
16         }
17         //可能有些行不存在平民，那么最终村在index + 1 < k的情况，需要按照行号从小到大存入返回vector
18         for (int i = 0; i < r; ++i) {
19             if (index < k && !visited[i]) {
20                 visited[i] = true;
21                 ans[index++] = i;
22             }
23         }
24         a:;
25         return ans;
26     }

python3代码：

1 class Solution:
2     def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
3         s = [(sum(m), i) for i, m in enumerate(mat)]
4         r = sorted(s)
5         return [r[1] for r in r[:k]]