LeetCode 1341. The K Weakest Rows in a Matrix

Question

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
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
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]

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]

Solution

We need to maintain the number of soldiers and each row, and then ascending by the count. Also, we can use the priority queue to reduce the time complexity for sorting because we just need to maintain k size. In addition, the question gives us the condition, which means all 1s show before than 0s, so we can use binary search for optimizing.

Code

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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
// time:O(nlogm) space:O(n), n = # of rows, m = # of cols.
public int[] kWeakestRows(int[][] mat, int k) {
if (mat == null || mat.length == 0 || mat[0] == null || mat[0].length == 0) return new int[]{};
PriorityQueue<Row> pq = new PriorityQueue<>(k, Collections.reverseOrder());
int m = mat.length;
int n = mat[0].length;
int[] res = new int[k];
for (int i = 0; i < m; i++) {
// use binary search to optimize
int count = find(mat[i], 1) + 1;
pq.offer(new Row(count, i));
if (pq.size() > k) pq.poll();
}
int i = k - 1;
while (!pq.isEmpty()) {
res[i--] = pq.poll().row;
}
return res;
}

public class Row implements Comparable<Row> {
int num;
int row;

public Row(int n, int r) {
num = n;
row = r;
}

@Override
public int compareTo(Row o) {
if (this.num != o.num) return this.num - o.num;
return this.row - o.row;
}
}

public int find(int[] nums, int target) {
int left = 0;
int right = nums.length - 1;
while (left + 1 < right) {
int mid = left + (right - left) / 2;
if (nums[mid] == target) left = mid;
else right = mid;
}
if (nums[right] == 1) return right;
if (nums[left] == 1) return left;
return -1;
}