1337.the-k-weakest-rows-in-a-matrix.md

1337. The K Weakest Rows in a Matrix

Heap

fun kWeakestRows(mat: Array<IntArray>, k: Int): IntArray {
    val count = IntArray(mat.size)
    val minHeap = PriorityQueue<Int>() { r1, r2 -> 
        if (count[r1] == count[r2]) r1 - r2
        else count[r1] - count[r2]
    }
    for (r in 0 until mat.size) {
        for (c in 0 until mat[0].size) {
            if (mat[r][c] == 1) count[r]++
            else break
        }
        minHeap.add(r)
    }
    val result = IntArray(k) { _ -> minHeap.poll() }
    return result
}

• Time Complexity: O(m * n * lg m).
• Space Complexity: O(m).

Binary Search

The 1s are always before 0s, so each row is sorted, and we can apply binary search to find the last 1 for counting (or the first 0).

1, 1, 0, 0, 0
   ^            // Last 1's
      *         // First 0's

fun kWeakestRows(mat: Array<IntArray>, k: Int): IntArray {
    val count = IntArray(mat.size)
    val maxHeap = PriorityQueue<Int>() { r1, r2 -> 
        if (count[r1] == count[r2]) r2 - r1
        else count[r2] - count[r1]
    }
    for (r in 0 until mat.size) {
        count[r] = findCount(mat[r])
        maxHeap.add(r)
        if (maxHeap.size > k) maxHeap.poll()
    }

    val result = IntArray(k)
    for (i in k - 1 downTo 0) {
        result[i] = maxHeap.poll()
    }
    return result
}

private fun findCount(A: IntArray): Int {
    var left = 0
    var right = A.size - 1
    while (left <= right) {
        val middle = left + (right - left) / 2
        if (A[middle] == 1) {
            left = middle + 1
        } else {
            right = middle - 1
        }
    }
    return right + 1
}

• Time Complexity: O(m * lg n + k * lg m).
• Space Complexity: O(m).