The K Weakest Rows in a Matrix

easy heap matrix

Problem

You are given an m × n binary matrix mat where 1 = soldier and 0 = civilian. Soldiers stand to the left of civilians in every row. Return the indices of the k weakest rows ordered from weakest to strongest — weaker means fewer soldiers, with row index as a tiebreaker.

Inputmat = [[1,1,0,0],[1,0,0,0],[1,1,1,1],[1,1,0,0]], k = 2

Output[1, 0]

Row strengths = [2,1,4,2]. Two weakest sorted by (count, idx): row 1 (1 soldier), then row 0 (2 soldiers, smaller index).

mat (JSON)

import heapq

def k_weakest_rows(mat, k):
    pairs = [(sum(row), i) for i, row in enumerate(mat)]
    return [i for _, i in heapq.nsmallest(k, pairs)]

function kWeakestRows(mat, k) {
  const pairs = mat.map(function (row, i) {
    return [row.reduce(function (a, b) { return a + b; }, 0), i];
  });
  pairs.sort(function (a, b) { return a[0] - b[0] || a[1] - b[1]; });
  return pairs.slice(0, k).map(function (p) { return p[1]; });
}

class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) ->
            a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);
        for (int i = 0; i < mat.length; i++) {
            int c = 0;
            for (int v : mat[i]) c += v;
            pq.offer(new int[]{c, i});
        }
        int[] ans = new int[k];
        for (int i = 0; i < k; i++) ans[i] = pq.poll()[1];
        return ans;
    }
}

vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
    vector<pair<int,int>> v;
    for (int i = 0; i < (int)mat.size(); i++) {
        int c = 0; for (int x : mat[i]) c += x;
        v.push_back({c, i});
    }
    sort(v.begin(), v.end());
    vector<int> ans;
    for (int i = 0; i < k; i++) ans.push_back(v[i].second);
    return ans;
}

Explanation

A row's "strength" is just how many soldiers (1s) it has, so the whole problem reduces to ranking rows by a number. We attach each row's soldier count to its index and then pick the k smallest.

For every row we build a pair (count, index) where count is the sum of that row (the soldier total). Pairing in (count, index) order is the trick that gives us the tiebreaker for free: when two rows have the same count, the smaller index naturally comes first.

We then ask for the k smallest pairs with heapq.nsmallest(k, pairs) and read off just the indices. Because tuples compare by count first and index second, the result is exactly weakest-to-strongest, ties broken by row order.

Example: mat = [[1,1,0,0],[1,0,0,0],[1,1,1,1],[1,1,0,0]], k = 2. The pairs are (2,0), (1,1), (4,2), (2,3). The two smallest are (1,1) and (2,0), so the answer is [1, 0].

Using nsmallest with a heap keeps this efficient even for tall matrices, since it tracks only the k best rows rather than fully sorting everything.

Time: O(m · n + m log k) Space: O(m)