題解

二分搜索+排序

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        m,n = len(mat),len(mat[0])
        def search_bound(nums: List[int]): # O(logn)
            nums = nums[::-1]
            n = len(nums)
            left, right = 0, n - 1
            bound = 0

            while left <= right:
                mid = left + (right - left) // 2
                if nums[mid] < 1:
                    left = mid + 1
                else:
                    bound = n - mid
                    right = mid - 1
            return bound
        
        output = []
        for row_index in range(m):
            total = search_bound(mat[row_index])
            output.append((total,row_index))
        output = sorted(output, key=lambda x: x[0]) # O(nlogn)
        return [output[i][-1] for i in range(k)]