Solution: The K Weakest Rows in a Matrix (ver. 2)

This is part of a series of Leetcode solution explanations (index). If you liked this solution or found it useful, please like this post and/or upvote my solution post on Leetcode's forums.

---

Note: This is my second version of a solution post for this problem. This version has a better time complexity (O(m * log(n + k)) vs. O(m * n)) as well as a better space complexity (O(k) vs. O(m)), but it does so by using a binary search function, a max-heap data structure, and bit manipulation, which are fairly complex tools for an "Easy" problem.

What's more is they don't provide any actual real-world performance increase when you consider the massive increase in process overhead and the relative small ranges of the inputs.

All told, I still prefer the first version because it's a simple approach more in line with an Easy problem.

After working through the more complicated code, however, I decided I might as well share it, as well as the reasoning behind it.

---

Leetcode Problem #1337 (Easy): The K Weakest Rows in a Matrix

---

Description:

Given a m * n matrix mat of ones (representing soldiers) and zeros (representing civilians), return the indexes of the k weakest rows in the matrix ordered from the weakest to the strongest.

A row i is weaker than row j, if the number of soldiers in row i is less than the number of soldiers in row j, or they have the same number of soldiers but i is less than j. Soldiers are always stand in the frontier of a row, that is, always ones may appear first and then zeros.

---

Examples:

Example 1:
Input:	mat = [[1,1,0,0,0],[1,1,1,1,0],[1,0,0,0,0],[1,1,0,0,0],[1,1,1,1,1]] k = 3
Output:	[2,0,3]
Explanation:	The number of soldiers for each row is: row 0 -> 2 row 1 -> 4 row 2 -> 1 row 3 -> 2 row 4 -> 5 Rows ordered from the weakest to the strongest are [2,0,3,1,4]

Example 2:
Input:	mat = [[1,0,0,0],[1,1,1,1],[1,0,0,0],[1,0,0,0]] k = 2
Output:	[0,2]
Explanation:	The number of soldiers for each row is: row 0 -> 1 row 1 -> 4 row 2 -> 1 row 3 -> 1 Rows ordered from the weakest to the strongest are [0,2,3,1]

---

Constraints:

• m == mat.length
• n == mat[i].length
• 2 <= n, m <= 100
• 1 <= k <= m
• mat[i][j] is either 0 or 1.

---

Idea:

In this approach, we can iterate through the rows and use the binary search function find the location of the first 0 in the row. Since we ultimately want our answer (ans) sorted first by earliest 0 and second by earliest index, we can go ahead and use that to our advantage by employing bit manipulation to condense the two values together before insertion into our max-heap (heap).

If we bitwise shift the location of the earliest 0 to the left and then add the row index, the resulting number should automatically sort precisely the way we need it to. A shift to the left by 7 should clear the max row index value of 99.

You can use either a min-heap or a max-heap, but a max-heap allows us very minor time and space complexity improvements by allowing us to keep the heap size down to k instead of n.

Once we're through iterating through the rows, we can just extract each value from heap, isolate just the row index with a bitwise AND, and insert it in reverse order into ans.

Eeeeasy peasy.

Unfortunately, while this is a more optimal solution on paper, it doesn't translate into actual benefit considering the value constraints implied on this Easy problem.

For the all four language examples below, it provided me pretty much identical real-world time and space results, with an awful lot more coding.

---

Implementation:

In all four languages, I used a custom binary search function, as the rows are in reverse order, so built-in functions like Python's bisect() and Java's Arrays.binarySearch() wouldn't work.

For Javascript, I used a custom max-heap implementation using a typed array for faster processing.

Python and Java both default to min-heap structures, so I just reversed the signs on the inputs to effectively simulate a max-heap structure.

---

Javascript Code:

var kWeakestRows = function(M, K) {
    let y = M.length, x = M[0].length,
        heap = new Int16Array(K+2), hix = 1,
         ans = new Uint8Array(K)
    heap[0] = 32767
    const heapify = val => {
        let i = hix, par = i >> 1, temp
        heap[hix++] = val
        while (heap[par] < heap[i]) {
            temp = heap[par], heap[par] = heap[i], heap[i] = temp
            i = par, par = i >> 1
        }
    }
    const extract = () => {
        let max = heap[1], left, right, temp,
            i = 1, child = heap[3] > heap[2] ? 3 : 2
        heap[1] = heap[--hix], heap[hix] = 0
        while (heap[i] < heap[child]) {
            temp = heap[child], heap[child] = heap[i], heap[i] = temp
            i = child, left = i << 1, right = left + 1
            child = heap[right] > heap[left] ? right : left
        }
        return max
    }
    const find = row => {
        let left = 0, right = x
        while (left <= right) {
            let mid = left + right >> 1
            if (row[mid] > 0) left = mid + 1
            else right = mid - 1
        }
        return left
    }
    for (let i = 0; i < y; i++) {
        heapify((find(M[i]) << 7) + i)
        if (hix > K + 1) extract()
    }
    while(hix) ans[hix-2] = extract() & (1 << 7) - 1
    return ans
};

---

Python Code:

class Solution:
    def kWeakestRows(self, M: List[List[int]], K: int) -> List[int]:
        y, x = len(M), len(M[0])
        def find(row: List[int]) -> int:
            left, right = 0, x
            while left <= right:
                mid = left + right >> 1
                if mid < x and row[mid] > 0: left = mid + 1
                else: right = mid - 1
            return left
        ans, heap = [0] * K, []
        for i in range(y):
            heappush(heap, (-find(M[i]) << 7) - i)
            if len(heap) > K: heappop(heap)
        while heap: ans[len(heap)] = -heappop(heap) & (1 << 7) - 1
        return ans

---

Java Code:

class Solution {
    public int[] kWeakestRows(int[][] M, int K) {
        int y = M.length, x = M[0].length;
        PriorityQueue<Integer> heap = new PriorityQueue<Integer>();
        int[] ans = new int[K];
        for (int i = 0; i < y; i++) {
            heap.add(-(find(M[i]) << 7) - i);
            if (heap.size() > K) heap.remove();
        }
        while (heap.size() > 0)
            ans[heap.size()-1] = -heap.remove() & (1 << 7) - 1;
        return ans;
    }
    int find(int[] row) {
        int x = row.length;
        int left = 0, right = x;
        while (left <= right) {
            int mid = left + right >> 1;
            if (mid < x && row[mid] > 0) left = mid + 1;
            else right = mid - 1;
        }
        return left;
    }
}

---

C++ Code:

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& M, int K) {
        int y = M.size(), x = M[0].size();
        vector<int> ans(K);
        priority_queue<int> heap;
        for (int i = 0; i < y; i++) {
            heap.push((find(M[i]) << 7) + i);
            if (heap.size() > K) heap.pop();
        }
        while (heap.size()) {
            ans[heap.size()-1] = heap.top() & (1 << 7) - 1;
            heap.pop();
        }
        return ans;
    }
    int find(vector<int> row) {
        int x = row.size();
        int left = 0, right = x;
        while (left <= right) {
            int mid = left + right >> 1;
            if (mid < x && row[mid] > 0) left = mid + 1;
            else right = mid - 1;
        }
        return left;
    }
};