The K Weakest Rows in a Matrix
EasyPrompt
You are given an m x n binary matrix mat of 1's (representing soldiers) and 0's (representing civilians). The soldiers are positioned in front of the civilians. That is, all the 1's will appear to the left of all the 0's in each row.
A row i is weaker than a row j if one of the following is true:
- The number of soldiers in row
iis less than the number of soldiers in rowj. - Both rows have the same number of soldiers and
i < j.
Return the indices of the k weakest rows in the matrix ordered from weakest to strongest.
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 in each row is:
- Row 0: 2
- Row 1: 4
- Row 2: 1
- Row 3: 2
- Row 4: 5
The rows ordered from weakest to 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 in each row is:
- Row 0: 1
- Row 1: 4
- Row 2: 1
- Row 3: 1
The rows ordered from weakest to strongest are [0,2,3,1].
Constraints:
m == mat.lengthn == mat[i].length2 <= n, m <= 1001 <= k <= mmatrix[i][j]is either 0 or 1.
Approaches
3 approaches with complexity analysis and trade-offs.
This approach involves two main steps. First, we iterate through the entire matrix to calculate the number of soldiers (strength) for each row. We store these strengths along with their original row indices. Second, we sort this collection based on the weakness criteria: primarily by strength in ascending order, and then by row index in ascending order for ties. Finally, we pick the first k elements from the sorted collection.
Algorithm
- Create a 2D array
strengthsof sizem x 2to store[soldier_count, row_index]. - Iterate from
i = 0tom-1:- Count the number of
1s inmat[i]by iterating through all its columns. - Store the count and the index
iinstrengths[i].
- Count the number of
- Sort the
strengthsarray. The primary sorting key is the soldier count (ascending), and the secondary key is the row index (ascending). - Create a result array
resultof sizek. - Populate
resultwith the indices from the firstkentries of the sortedstrengthsarray. - Return
result.
Walkthrough
The most straightforward way to solve the problem is to first determine the strength of each row, then sort the rows based on these strengths, and finally pick the top k.
- Strength Calculation: We create an auxiliary 2D array or a list of pairs to hold the strength and original index of each row. We loop through each row of the input matrix
mat. For each row, we perform a linear scan to count the number of1s. This count is the row's strength. - Sorting: After calculating all strengths, we sort the auxiliary array. The sorting logic must follow the problem's definition of weakness: if two rows have different strengths, the one with fewer soldiers comes first. If they have the same strength, the one with the smaller original index comes first.
- Result Extraction: Once sorted, the first
kelements in our auxiliary array correspond to thekweakest rows. We create a result array of sizekand fill it with the original indices from thesekelements.
import java.util.Arrays; class Solution { public int[] kWeakestRows(int[][] mat, int k) { int m = mat.length; int[][] strengths = new int[m][2]; for (int i = 0; i < m; i++) { int count = 0; for (int val : mat[i]) { if (val == 1) { count++; } else { break; // Optimization: since 1s are followed by 0s } } strengths[i][0] = count; strengths[i][1] = i; } Arrays.sort(strengths, (a, b) -> { if (a[0] != b[0]) { return a[0] - b[0]; } else { return a[1] - b[1]; } }); int[] result = new int[k]; for (int i = 0; i < k; i++) { result[i] = strengths[i][1]; } return result; }}Complexity
Time
O(m * n + m * log(m)) - Calculating strengths for all `m` rows takes `O(m * n)`. Sorting `m` items takes `O(m * log(m))`. The total complexity is dominated by the larger of these two terms.
Space
O(m) - We need to store the strength and index for all `m` rows in an auxiliary array.
Trade-offs
Pros
Simple to understand and implement.
Cons
Inefficient for wide matrices (large
n) as it scans every element, leading to a higher time complexity.Requires sorting all
mrows, which is unnecessary work ifkis much smaller thanm.
This approach improves upon the first one by optimizing the strength calculation. Since each row is sorted with 1s followed by 0s, we can use binary search to find the number of soldiers instead of a linear scan. The rest of the logic, which involves sorting all m rows and picking the top k, remains the same.
Algorithm
- Create a 2D array
strengthsof sizem x 2. - Iterate from
i = 0tom-1:- Use binary search on
mat[i]to find the number of1s (soldiers). This is possible because all1s appear before0s. - Store the count and the index
iinstrengths[i].
- Use binary search on
- Sort the
strengthsarray based on the count, then by index for ties. - Create a result array
resultof sizek. - Populate
resultwith the firstkindices from the sortedstrengthsarray. - Return
result.
Walkthrough
We can make the strength calculation more efficient. The problem statement guarantees that in each row, all 1s appear before all 0s. This sorted property allows us to use binary search to find the number of soldiers.
- Optimized Strength Calculation: We iterate through each row
i. For each row, we perform a binary search to find the index of the first0. This index is equivalent to the number of1s (soldiers) in that row. If a row contains all1s, the binary search will conclude that the count isn. - Storage and Sorting: As in the previous approach, we store these
[strength, index]pairs in an auxiliary array. Then, we sort this array using the same comparison logic (strength first, then index). - Result Extraction: Finally, we extract the indices of the first
kelements from the sorted array.
This method significantly reduces the time spent on counting soldiers from O(n) to O(log n) per row.
import java.util.Arrays; class Solution { public int[] kWeakestRows(int[][] mat, int k) { int m = mat.length; int n = mat[0].length; int[][] strengths = new int[m][2]; for (int i = 0; i < m; i++) { strengths[i][0] = countSoldiers(mat[i], n); strengths[i][1] = i; } Arrays.sort(strengths, (a, b) -> { if (a[0] != b[0]) { return a[0] - b[0]; } else { return a[1] - b[1]; } }); int[] result = new int[k]; for (int i = 0; i < k; i++) { result[i] = strengths[i][1]; } return result; } // Binary search to find the count of soldiers (1s) private int countSoldiers(int[] row, int n) { int low = 0; int high = n - 1; while (low <= high) { int mid = low + (high - low) / 2; if (row[mid] == 1) { low = mid + 1; } else { high = mid - 1; } } return low; }}Complexity
Time
O(m * log(n) + m * log(m)) - For each of the `m` rows, binary search takes `O(log n)`. Sorting `m` items takes `O(m * log(m))`.
Space
O(m) - Space is required to store the strength and index for all `m` rows before sorting.
Trade-offs
Pros
Faster strength calculation (
O(log n)) compared to the brute-force approach (O(n)), making it more efficient for wide matrices.
Cons
Still requires sorting all
mrows, which is inefficient ifkis much smaller thanm.Uses
O(m)space, which can be suboptimal ifmis large andkis small.
This is the most efficient approach, especially when k is much smaller than m. It combines the optimized strength calculation from the second approach with a max-heap to find the k weakest rows without sorting the entire set. We maintain a max-heap of size k. For each row, we calculate its strength and add it to the heap. If the heap size exceeds k, we remove the strongest element, ensuring the heap always contains the k weakest rows seen so far.
Algorithm
- Initialize a max-heap (Priority Queue) of size
k. The comparator should prioritize higher strength, then higher index, to keep the 'strongest' elements at the top for easy removal. - Iterate from
i = 0tom-1:- Calculate the strength of
mat[i]using binary search. - Offer the
[strength, i]pair to the heap. - If
heap.size() > k,poll()the heap to remove the current strongest element.
- Calculate the strength of
- After the loop, the heap contains the
kweakest rows. - Create a result array
resultof sizek. - Populate
resultfrom indexk-1down to0by repeatedly polling from the heap. This reverses the max-heap order to the required weakest-to-strongest order. - Return
result.
Walkthrough
This approach optimizes both time and space by avoiding a full sort and using a data structure tailored for 'Top K' problems: a max-heap (or a PriorityQueue in Java).
- Max-Heap Setup: We initialize a max-heap that will store pairs of
[strength, index]. The heap's comparator is crucial: it must treat rows with more soldiers as 'greater'. For ties in strength, the row with the larger index is 'greater'. This ensures that when we remove an element, we are removing the strongest row according to the problem's rules. - Processing Rows: We iterate through each row of the matrix. For each row, we calculate its strength using binary search (
O(log n)). We then add this[strength, index]pair to our max-heap. Immediately after adding, we check if the heap's size has exceededk. If it has, wepoll()the heap, which removes the maximum (strongest) element. This way, the heap's size is always maintained at or belowk, and it always holds thekweakest rows encountered so far. - Result Extraction: After processing all
mrows, the max-heap contains thekweakest rows. However, since it's a max-heap, polling will return them from strongest to weakest. To get the desired output order, we create a result array of sizekand fill it from back to front (from indexk-1down to0) by repeatedly polling from the heap.
import java.util.PriorityQueue; class Solution { public int[] kWeakestRows(int[][] mat, int k) { int m = mat.length; int n = mat[0].length; // Max-heap to store pairs [strength, index] // Strongest rows will have higher priority and will be polled first if size > k PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> { if (a[0] != b[0]) { return b[0] - a[0]; // Sort by strength descending } else { return b[1] - a[1]; // Sort by index descending } }); for (int i = 0; i < m; i++) { int strength = countSoldiers(mat[i], n); maxHeap.offer(new int[]{strength, i}); if (maxHeap.size() > k) { maxHeap.poll(); } } int[] result = new int[k]; for (int i = k - 1; i >= 0; i--) { result[i] = maxHeap.poll()[1]; } return result; } private int countSoldiers(int[] row, int n) { int low = 0, high = n - 1; while (low <= high) { int mid = low + (high - low) / 2; if (row[mid] == 1) { low = mid + 1; } else { high = mid - 1; } } return low; }}Complexity
Time
O(m * log(n) + m * log(k)) - For each of the `m` rows, we perform a binary search (`O(log n)`) and a heap operation (`O(log k)`). The final extraction from the heap takes `O(k * log k)`, which is subsumed by the main loop's complexity.
Space
O(k) - The max-heap stores at most `k` elements.
Trade-offs
Pros
Most efficient in both time and space.
Avoids sorting all
melements, making it much faster whenkis small compared tom.Space complexity is proportional to
k, notm, which is a significant improvement for largem.
Cons
Slightly more complex to implement due to the custom heap comparator and the final extraction step which requires careful handling to get the correct order.
Solutions
Solution
class Solution {public int[] kWeakestRows(int[][] mat, int k) { int m = mat.length, n = mat[0].length; int[] res = new int[m]; List<Integer> idx = new ArrayList<>(); for (int i = 0; i < m; ++i) { idx.add(i); int[] row = mat[i]; int left = 0, right = n; while (left < right) { int mid = (left + right) >> 1; if (row[mid] == 0) { right = mid; } else { left = mid + 1; } } res[i] = left; } idx.sort(Comparator.comparingInt(a->res[a])); int[] ans = new int[k]; for (int i = 0; i < k; ++i) { ans[i] = idx.get(i); } return ans; }}class Solution: def kWeakestRows(self, mat: List[List[int]], k: int) -> List[int]: m, n = len(mat), len(mat[0]) ans = [n - bisect_right(row[:: - 1], 0) for row in mat] idx = list(range(m)) idx . sort(key=lambda i: ans[i]) return idx[: k]class Solution { public: int search ( vector < int >& m ) { int l = 0 ; int h = m . size () - 1 ; while ( l <= h ) { int mid = l + ( h - l ) / 2 ; if ( m [ mid ] == 0 ) h = mid - 1 ; else l = mid + 1 ; } return l ; } vector < int > kWeakestRows ( vector < vector < int >>& mat , int k ) { vector < pair < int , int >> p ; vector < int > res ; for ( int i = 0 ; i < mat . size (); i ++ ) { int count = search ( mat [ i ]); p . push_back ({ count , i }); } sort ( p . begin (), p . end ()); for ( int i = 0 ; i < k ; i ++ ) { res . push_back ( p [ i ]. second ); } return res ; } };Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.