文章目录
- 题目描述
- 题目难度——简单
- 方法一:暴力,统计
- 代码/Python
- 方法二:优化
- 代码
- 总结
- 彩蛋
题目描述
给你一个大小为 m * n 的矩阵 mat,矩阵由若干军人和平民组成,分别用 1 和 0 表示。
请你返回矩阵中战斗力最弱的 k 行的索引,按从最弱到最强排序。
如果第 i 行的军人数量少于第 j 行,或者两行军人数量相同但 i 小于 j,那么我们认为第 i 行的战斗力比第 j 行弱。
军人 总是 排在一行中的靠前位置,也就是说 1 总是出现在 0 之前。
示例 1:
输入: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
输出:[2,0,3]
解释: 每行中的军人数目: 行 0 -> 2 行 1 -> 4 行 2> -> 1 行 3 -> 2 行 4 -> 5 从最弱到最强对这些行排序后得到 [2,0,3,1,4]
示例 2:
输入:mat = [[1,0,0,0], [1,1,1,1], [1,0,0,0], [1,0,0,0]], k = 2
输出:[0,2]
解释: 每行中的军人数目: 行 0 -> 1 行 1 -> 4 行 2 -> 1 行 3 -> 1
从最弱到最强对这些行排序后得到 [0,2,3,1]
提示:
- m == mat.length
- n == mat[i].length
- 2 <= n, m <= 100
- 1 <= k <= m
- matrix[i][j] 不是 0 就是 1
题目链接
题目难度——简单
方法一:暴力,统计
题目的数据量很小,因此可以暴力的遍历,用一个数组保存每行的战斗力,最后把战斗力和下标组成一起,排序,返回前k个就行。
代码/Python
class Solution:
def kWeakestRows(self, mat: list[list[int]], k: int) -> list[int]:
m, n = len(mat), len(mat[0])
soldiers = [sum(row) for row in mat]
fights = [(i, x) for i, x in enumerate(soldiers)]
fights.sort(key = lambda x: x[1])
return [fights[i][0] for i in range(k)]
方法二:优化
注意题目中说的,每行中1都在前面,所以我们可以利用这个特点,我们可以不计数每行的战斗力,取而代之,获取每行最后一个1的下标也可以作为战斗力的指标。可以用二分查找来进行这一步。
代码
class Solution:
def kWeakestRows(self, mat: list[list[int]], k: int ) -> list[int]:
# 方法二,利用上上面那个性质,用二分查找
m, n = len(mat), len(mat[0])
soldiers = [0] * m
for i in range(m):
l, r, pos = 0, n - 1, -1
while l <= r:
mid = (l + r) // 2
if mat[i][mid] == 0:
r = mid - 1
else:
pos = mid
l = mid + 1
soldiers[i] = pos
# print(soldiers)
fight = sorted([(i, x) for i, x in enumerate(soldiers)], key = lambda x: x[1])
return [fight[i][0] for i in range(k)]
这么慢,挺离谱的,排前面的也有这种方法。后面排序可以用小根堆来代替,在前面获得每行的战斗力之后,将战斗力数组用一个小根堆来存储,取k个元素就得到了答案,不过还是得排个序,因为题目要求有序。
总结
方法一需要O(N2)的复杂度,方法二用二分查找优化后复杂度可以到O(NlogN),两种的空间都是O(N)。
彩蛋
这道题是去年8月1日的每日一题,题目中每行都是表示军人的1在前,致敬了我们的解放军,建军节快乐。