1337. The K Weakest Rows in a Matrix

https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/

Python

Min Heap

import heapq


class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        heap = []
        for row in mat:
            total = 0
            for num in row:
                if num == 1:
                    total += 1
                else:
                    break
            heap.append(total)

        return [
            index for index, _ in
            heapq.nsmallest(
                k, enumerate(heap), key=lambda item: item[1]
            )
        ]