[LeetCode] 1337. The K Weakest Rows in a Matrix

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.

矩阵中战斗力最弱的 K 行。

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

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

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

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/the-k-weakest-rows-in-a-matrix
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

这道题给的tag是二分法。二分法是用来快速找到每一行中 1 的出现次数。因为是需要找战斗力最弱的前 K 行，所以我们需要一个优先队列，存储K行。优先队列的规则是战斗力更大的在堆顶，或者战斗力相同的时候，行index更大的在堆顶。我们遍历所有的行并且用helper函数计算了所有行的战斗力，再经过优先队列刷一遍之后，留下的K行就是战斗力最弱的。同时注意因为堆顶那一行的战斗力是最强的所以写入结果集的时候要逆序。

时间O(mlogn * logk) - 有M行，每行都做二分法并被pq排序

空间O(n * k) - 有K行会被存到结果集，每行N个元素

Java实现

 1 class Solution {
 2     public int[] kWeakestRows(int[][] mat, int k) {
 3         // 战斗力大的在前，或行index大的在前
 4         PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> a[0] != b[0] ? b[0] - a[0] : b[1] - a[1]);
 5         int[] res = new int[k];
 6 
 7         for (int i = 0; i < mat.length; i++) {
 8             queue.offer(new int[] {helper(mat[i]), i});
 9             if (queue.size() > k) {
10                 queue.poll();
11             }
12         }
13 
14         while (k > 0) {
15             res[--k] = queue.poll()[1];
16         }
17         return res;
18     }
19 
20     private int helper(int[] row) {
21         int lo = 0;
22         int hi = row.length;
23         while (lo < hi) {
24             int mid = lo + (hi - lo) / 2;
25             if (row[mid] == 1) {
26                 lo = mid + 1;
27             } else {
28                 hi = mid;
29             }
30         }
31         return lo;
32     }
33 }

LeetCode 题目总结