1337. The K Weakest Rows in a Matrix #

Problem #

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.

Problem Summary #

You are given a matrix mat of size m * n, consisting of soldiers and civilians, represented by 1 and 0 respectively. Return the indexes of the k weakest rows in the matrix, sorted from weakest to strongest. If row i has fewer soldiers than row j, or the two rows have the same number of soldiers but i is less than j, then we consider row i weaker than row j. Soldiers are always positioned at the front of a row, meaning 1 always appears before 0.

Solution Ideas #

Easy problem. The first solution that comes to mind is to count the number of 1s in each row, then sort the results by the number of 1s in ascending order; if the number of 1s is the same, sort by row index in ascending order. Take the first K elements from the sorted array as the answer.
There is also a second solution to this problem. In the first solution, the condition in the problem statement, “Soldiers are always positioned at the front of a row, meaning 1 always appears before 0,” is not used. Because of this condition, if we traverse by columns, the rows where 0 appears first are the weakest rows. Rows with smaller indices are visited first, so among rows with the same number of 1s, the row with the smaller index will come first. Finally, remember to add the rows that are all 1s. Similarly, take the first K elements of the final output as the answer. The second solution is the most elegant and efficient one for this problem.

Code #

package leetcode

func kWeakestRows(mat [][]int, k int) []int {
	res := []int{}
	for j := 0; j < len(mat[0]); j++ {
		for i := 0; i < len(mat); i++ {
			if mat[i][j] == 0 && ((j == 0) || (mat[i][j-1] != 0)) {
				res = append(res, i)
			}
		}
	}
	for i := 0; i < len(mat); i++ {
		if mat[i][len(mat[0])-1] == 1 {
			res = append(res, i)
		}
	}
	return res[:k]
}