The K Weakest Rows in a Matrix

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        priority_queue<pair<int,int>> pq;
        for (int i = 0; i < (int)mat.size(); ++i) {
            int cnt = lower_bound(mat[i].begin(), mat[i].end(), 0, greater<int>()) - mat[i].begin();
            pq.push({cnt, i});
            if ((int)pq.size() > k) pq.pop();
        }
        vector<pair<int,int>> kept;
        while (!pq.empty()) { kept.push_back(pq.top()); pq.pop(); }
        sort(kept.begin(), kept.end());
        vector<int> ans;
        for (auto& [c, i] : kept) ans.push_back(i);
        return ans;
    }
};

import (
  "container/heap"
  "sort"
)

type rowItem struct{ cnt, idx int }
type rowMaxHeap []rowItem

func (h rowMaxHeap) Len() int { return len(h) }
func (h rowMaxHeap) Less(i, j int) bool {
  if h[i].cnt != h[j].cnt {
    return h[i].cnt > h[j].cnt
  }
  return h[i].idx > h[j].idx
}
func (h rowMaxHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *rowMaxHeap) Push(x interface{}) { *h = append(*h, x.(rowItem)) }
func (h *rowMaxHeap) Pop() interface{}   { o := *h; n := len(o); x := o[n-1]; *h = o[:n-1]; return x }

func kWeakestRows(mat [][]int, k int) []int {
  pq := &rowMaxHeap{}
  for i, row := range mat {
    cnt := sort.Search(len(row), func(j int) bool { return row[j] == 0 })
    heap.Push(pq, rowItem{cnt, i})
    if pq.Len() > k {
      heap.Pop(pq)
    }
  }
  kept := make([]rowItem, 0, k)
  for pq.Len() > 0 {
    kept = append(kept, heap.Pop(pq).(rowItem))
  }
  sort.Slice(kept, func(i, j int) bool {
    if kept[i].cnt != kept[j].cnt {
      return kept[i].cnt < kept[j].cnt
    }
    return kept[i].idx < kept[j].idx
  })
  ans := make([]int, 0, len(kept))
  for _, r := range kept {
    ans = append(ans, r.idx)
  }
  return ans
}

import bisect
import heapq
from typing import List

class Solution:
    def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]:
        # rows are 1s then 0s; count of 1s = bisect_left on reversed
        heap: List[tuple] = []
        for i, row in enumerate(mat):
            # find first 0; bisect_right against 1 in descending sequence
            # equivalent: count = sum(row), but binary search is O(log n)
            lo, hi = 0, len(row)
            while lo < hi:
                mid = (lo + hi) // 2
                if row[mid] == 1:
                    lo = mid + 1
                else:
                    hi = mid
            cnt = lo
            # max-heap by (-cnt, -i) to keep weakest k
            heapq.heappush(heap, (-cnt, -i))
            if len(heap) > k:
                heapq.heappop(heap)
        kept = [(-c, -i) for c, i in heap]
        kept.sort()
        return [i for _, i in kept]

class MaxHeap {
    constructor(cmp) { this.h = []; this.cmp = cmp; }
    size() { return this.h.length; }
    top() { return this.h[0]; }
    push(x) {
        this.h.push(x);
        let i = this.h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.cmp(this.h[p], this.h[i]) >= 0) break;
            [this.h[p], this.h[i]] = [this.h[i], this.h[p]];
            i = p;
        }
    }
    pop() {
        const top = this.h[0];
        const last = this.h.pop();
        if (this.h.length) {
            this.h[0] = last;
            let i = 0;
            const n = this.h.length;
            while (true) {
                let l = i * 2 + 1, r = i * 2 + 2, s = i;
                if (l < n && this.cmp(this.h[l], this.h[s]) > 0) s = l;
                if (r < n && this.cmp(this.h[r], this.h[s]) > 0) s = r;
                if (s === i) break;
                [this.h[s], this.h[i]] = [this.h[i], this.h[s]];
                i = s;
            }
        }
        return top;
    }
}

var kWeakestRows = function(mat, k) {
    // max-heap by (cnt, idx) — weakest stays at top after pop overflow
    const pq = new MaxHeap((a, b) => a[0] !== b[0] ? a[0] - b[0] : a[1] - b[1]);
    for (let i = 0; i < mat.length; i++) {
        const row = mat[i];
        let lo = 0, hi = row.length;
        while (lo < hi) {
            const mid = (lo + hi) >> 1;
            if (row[mid] === 1) lo = mid + 1;
            else hi = mid;
        }
        pq.push([lo, i]);
        if (pq.size() > k) pq.pop();
    }
    const kept = [];
    while (pq.size()) kept.push(pq.pop());
    kept.sort((a, b) => a[0] !== b[0] ? a[0] - b[0] : a[1] - b[1]);
    return kept.map(p => p[1]);
};

class Solution {
    public int[] kWeakestRows(int[][] mat, int k) {
        // max-heap by (cnt, idx) — weakest stays after popping overflow
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) ->
            a[0] != b[0] ? b[0] - a[0] : b[1] - a[1]);
        for (int i = 0; i < mat.length; i++) {
            int[] row = mat[i];
            int lo = 0, hi = row.length;
            while (lo < hi) {
                int mid = (lo + hi) >>> 1;
                if (row[mid] == 1) lo = mid + 1;
                else hi = mid;
            }
            pq.offer(new int[]{lo, i});
            if (pq.size() > k) pq.poll();
        }
        List<int[]> kept = new ArrayList<>();
        while (!pq.isEmpty()) kept.add(pq.poll());
        kept.sort((a, b) -> a[0] != b[0] ? a[0] - b[0] : a[1] - b[1]);
        int[] ans = new int[kept.size()];
        for (int i = 0; i < kept.size(); i++) ans[i] = kept.get(i)[1];
        return ans;
    }
}

class MaxHeap {
    h: number[][] = [];
    cmp: (a: number[], b: number[]) => number;
    constructor(cmp: (a: number[], b: number[]) => number) { this.cmp = cmp; }
    size(): number { return this.h.length; }
    push(x: number[]): void {
        this.h.push(x);
        let i: number = this.h.length - 1;
        while (i > 0) {
            const p: number = (i - 1) >> 1;
            if (this.cmp(this.h[p], this.h[i]) >= 0) break;
            [this.h[p], this.h[i]] = [this.h[i], this.h[p]];
            i = p;
        }
    }
    pop(): number[] {
        const top: number[] = this.h[0];
        const last: number[] = this.h.pop() as number[];
        if (this.h.length) {
            this.h[0] = last;
            let i: number = 0;
            const n: number = this.h.length;
            while (true) {
                let l: number = i * 2 + 1, r: number = i * 2 + 2, s: number = i;
                if (l < n && this.cmp(this.h[l], this.h[s]) > 0) s = l;
                if (r < n && this.cmp(this.h[r], this.h[s]) > 0) s = r;
                if (s === i) break;
                [this.h[s], this.h[i]] = [this.h[i], this.h[s]];
                i = s;
            }
        }
        return top;
    }
}

function kWeakestRows(mat: number[][], k: number): number[] {
    const pq: MaxHeap = new MaxHeap((a: number[], b: number[]) =>
        a[0] !== b[0] ? a[0] - b[0] : a[1] - b[1]);
    for (let i = 0; i < mat.length; i++) {
        const row: number[] = mat[i];
        let lo: number = 0, hi: number = row.length;
        while (lo < hi) {
            const mid: number = (lo + hi) >> 1;
            if (row[mid] === 1) lo = mid + 1;
            else hi = mid;
        }
        pq.push([lo, i]);
        if (pq.size() > k) pq.pop();
    }
    const kept: number[][] = [];
    while (pq.size()) kept.push(pq.pop());
    kept.sort((a: number[], b: number[]) => a[0] !== b[0] ? a[0] - b[0] : a[1] - b[1]);
    return kept.map((p: number[]) => p[1]);
}