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

请你返回方阵中战斗力最弱的 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

解题思路

最简单的思路就是采用排序（针对每行1的个数）。

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        return sorted(range(len(mat)), key=lambda x: sum(mat[x]))[:k]

这个问题和Leetcode 347:前K个高频元素同属于Top K问题，采用堆结构处理。

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        q = []
        for i, r in enumerate(mat):
            heapq.heappush(q, (-sum(r), -i))
            if len(q) > k:
                heapq.heappop(q)
        return [-t[1] for t in heapq.nlargest(k, q)]

这个问题中的一个条件我们一直没有使用"1 总是出现在 0 之前"，因此我们可以使用二分法统计1出现的次数。

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        def numOnes(row):
            l, r = 0, len(row)
            while l < r:
                mid = (l + r) >> 1
                if row[mid] == 1:
                    l = mid + 1
                else:
                    r = mid
            return l
       	
        q = []
        for i, r in enumerate(mat):
            heapq.heappush(q, (-numOnes(r), -i))
            if len(q) > k:
                heapq.heappop(q)
        return [-t[1] for t in heapq.nlargest(k, q)]

我将该问题的其他语言版本添加到了我的GitHub Leetcode

如有问题，希望大家指出！！！