Problem

1337. The K Weakest Rows in a Matrix

Solution

• Uses Binary search to count 1s
• Uses heap to keep k weakest rows
• It naturally satisfies the condition for weak rows with the same number of soldiers. The smaller index row is processed first and hence any next row would not get pushed to the heap.

from heapq import heappushpop, heappush, heappop

class Solution(object):
    def countOne(self, row):
        ans = 0
        for n in row:
            if n == 0:
                break
            ans += 1
        return ans
    
    def countOneBS(self, row, start, end):
        if row[start] == 0:
            return 0
        
        if row[end] == 1:
            return end - start + 1
        
        mid = start + (end - start) // 2
        
        # print("mid=", mid, "start=", start, "end=", end)
        return self.countOneBS(row, start, mid) + \
               self.countOneBS(row, mid + 1, end)
    
    
    def kWeakestRows(self, mat, k):
        """
        :type mat: List[List[int]]
        :type k: int
        :rtype: List[int]
        """
        # max heap
        heap = []
        weakestRows = []
        
        for index, row in enumerate(mat):
            c = self.countOneBS(row, 0, len(row) - 1)
            
            if len(heap) == k:
                heappushpop(heap, (-c, -index))
            else:
                heappush(heap,(-c, -index))
                
        while heap:
            weakestRows.append(-heappop(heap)[1])
            
        return weakestRows[::-1]

Reference

• https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/
• https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/discuss/496713/Python-One-Liner-using-Sorting