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

What we can do is to count number of ones in each row and then sort all rows, using key with two elements: (sum(row), ind), because if we have two rows with equal number of ones, we need to choose the one with smallest index.

Complexity: time complexity is O(mn) to evaluate all sums and O(m*log m) to perform sort and choose k smallest elements, so final complexity is O(mn + m log m). Space complexity is O(m).

class Solution:
    def kWeakestRows(self, mat, k):
        return [j for _, j in sorted([(sum(row), ind) for ind, row in enumerate(mat)])[:k]]

Buckets solution

This problem is marked as easy on leetcode, because dimensions of problem is quite small and given solution will be fast enough. However we did not use some information given in problem statement, so we can improve our performance:

  1. First of all, in each row data is sorted, like [1, 1, 1, ..., 0, 0, 0], and we do not really need to traverse all list to find how many ones we have inside, and we can do binary search instead with complexity O(log n), so complexity of first step will be O(m log n) instead of O(mn). Unfortunatelly you can not use bisect library here, because data in rows sorted in reverse order and you need to do it by hands.
  2. Secondly, we do not need to sort all data, but we only need to choose k smallest elements, which can be done, using heaps: we can keep heap with k elements and have complexity O(m log k). Moreover, because frequencies will not be more than n, we can use bucket sort to find k smallest values.

Final time complexity will be O(m log n + n), space complexity is O(n). If m log k << n, it is more preferable to use heaps, however even though complexity of heap operations is O(log k), constant is quite big, so I can say, bucket sort is preferable.

class Solution:
    def kWeakestRows(self, mat, k):
        def num_ones(L):
            beg, end = 0, n
            while beg < end:
                mid = (beg + end)//2
                if L[mid] > 0: 
                    beg = mid + 1
                else: 
                    end = mid
            return beg
        
        m, n = len(mat), len(mat[0])
        buckets = [[] for _ in range(n+1)]
        for i, row in enumerate(mat):
            buckets[num_ones(row)].append(i)
            
        return list(chain(*buckets))[:k]

If you like the solution, you can upvote it on leetcode discussion section: Problem 1337