You are given an m x n
binary matrix mat
of 1
's (representing soldiers) and 0
's (representing civilians). The soldiers are positioned in front of the civilians. That is, all the 1
's will appear to the left of all the 0
's in each row.
A row i
is weaker than a row j
if one of the following is true:
i
is less than the number of soldiers in row j
.i < j
.Return the indices of the k
weakest rows in the matrix ordered from weakest to strongest.
For example:
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]
Explain your approach and provide code with its time and space complexity. Also discuss different edge cases that can occur with this problem and how to handle them.
We can iterate through each row, count the number of soldiers (1s), and store the row index along with the count. Then, we sort these pairs based on the soldier count (and row index as a tie-breaker). Finally, we extract the first k
row indices from the sorted list.
def k_weakest_rows_naive(mat, k):
rows = []
for i, row in enumerate(mat):
count = sum(row)
rows.append((count, i))
rows.sort()
result = [row[1] for row in rows[:k]]
return result
To optimize the soldier counting process, we can use binary search. Since the soldiers are always in front, we can binary search for the index of the first civilian (0) in each row. This gives us the count of soldiers in that row. Then the rest of the algorithm is the same as in the naive solution, but with improved time complexity for soldier counting.
def k_weakest_rows_optimal(mat, k):
def count_soldiers(row):
left, right = 0, len(row)
while left < right:
mid = (left + right) // 2
if row[mid] == 0:
right = mid
else:
left = mid + 1
return left
rows = []
for i, row in enumerate(mat):
count = count_soldiers(row)
rows.append((count, i))
rows.sort()
result = [row[1] for row in rows[:k]]
return result
In cases where n > m, the optimal solution provides a significant improvement in runtime efficiency.
Both solutions have the same space complexity because they both store information for each row.
Here's the updated code to handle these edge cases:
def k_weakest_rows_optimal_edge_cases(mat, k):
if not mat:
return []
def count_soldiers(row):
if not row:
return 0
left, right = 0, len(row)
while left < right:
mid = (left + right) // 2
if row[mid] == 0:
right = mid
else:
left = mid + 1
return left
rows = []
for i, row in enumerate(mat):
count = count_soldiers(row)
rows.append((count, i))
rows.sort()
result = [row[1] for row in rows[:min(k, len(mat))]]
return result