Find Valid Matrix Given Row and Column Sums Solution In C++/Python/java/JS
Problem Description
You are given two arrays rowSum and colSum of non-negative integers where rowSum[i] is the sum of the elements in the ith row and colSum[j] is the sum of the elements of the jth column of a 2D matrix. In other words, you do not know the elements of the matrix, but you do know the sums of each row and column.
Find any matrix of non-negative integers of size rowSum.length x colSum.length that satisfies the rowSum and colSum requirements.
Return a 2D array representing any matrix that fulfills the requirements. It's guaranteed that at least one matrix that fulfills the requirements exists.
Example 1:
Input: rowSum = [3,8], colSum = [4,7]
Output: [[3,0], [1,7]]
Explanation:
0th row: 3 + 0 = 3 == rowSum[0]
1st row: 1 + 7 = 8 == rowSum[1]
0th column: 3 + 1 = 4 == colSum[0]
1st column: 0 + 7 = 7 == colSum[1]
The row and column sums match, and all matrix elements are non-negative.
Another possible matrix is: [[1,2], [3,5]]
Example 2:
Input: rowSum = [5,7,10], colSum = [8,6,8]
Output: [[0,5,0],
[6,1,0],
[2,0,8]]
Constraints
- 1 <= rowSum.length, colSum.length <= 500
- 0 <= rowSum[i], colSum[i] <= 10^8
- sum(rowSum) == sum(colSum)
Finding the k-th largest element in an array is a well-known problem in computer science, and several efficient techniques have been developed to solve it. Let's explore different approaches to tackle this problem effectively
Brute Force Approach
Intuition
The goal is to find the indices of the k weakest rows in a given matrix. Each row consists of soldiers (represented by 1s) followed by civilians (represented by 0s). A row is considered weaker if it has fewer soldiers than another row, and in case of a tie, the row appearing earlier in the matrix is weaker.
To achieve this, the first step is to determine the strength of each row. Since the soldiers always appear before the civilians, counting the number of 1s in a row directly gives us its strength. We then store this count along with the row index, as the index is needed for tie-breaking when two rows have the same number of soldiers.
Once we have collected the strength of all rows, the next step is to sort them. Sorting ensures that the rows with fewer soldiers appear first. If two rows have the same number of soldiers, the row with the smaller index comes first due to the natural ordering. This guarantees that we can extract the k weakest rows efficiently.
Finally, after sorting, we extract the indices of the first k rows, as these represent the weakest rows in the matrix. Sorting takes O(m log m) time, where m is the number of rows, and counting soldiers for each row takes O(n), making the overall approach efficient for moderate input sizes.
Approach
Step 1: Count the Soldiers in Each Row
- Iterate through each row of the matrix and count the number of soldiers (1s) present.
- Since soldiers always appear before civilians, counting the 1s directly gives the row’s strength.
- Store this soldier count along with the row index to maintain the original position of each row.
Step 2: Sort Rows by Strength
- Sort the rows based on the number of soldiers in ascending order, ensuring weaker rows come first.
- If two rows have the same number of soldiers, sort them by their original index to maintain the correct order.
Step 3: Extract the K Weakest Rows
- After sorting, select the first k rows from the sorted list.
- Retrieve only their indices to form the final output.
Step 3: Return the Result
- Return the list of k row indices representing the weakest rows in the matrix.
Dry Run
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
Step 1: Count the Soldiers in Each Row
Count the number of 1s in each row:
- Row 0: 2 soldiers
- Row 1: 4 soldiers
- Row 2: 1 soldier
- Row 3: 2 soldiers
- Row 4: 5 soldiers
Store the soldier count along with row indices: [(2,0),(4,1),(1,2),(2,3),(5,4)]
Step 2: Sort Rows by Strength
- Sort the rows based on soldier count (ascending order). If two rows have the same count, sort by row index:
- Sorted order: [(1,2),(2,0),(2,3),(4,1),(5,4)]
- Ordered row indices: [2,0,3,1,4]
Step 3: Extract the K Weakest Rows
- Select the first k = 3 weakest rows: [2,0,3]
Step 4: Return the Result
- Output: [2,0,3]
Code for All Languages
C++
// Approach: Sorting-Based Solution to Find the K Weakest Rows in a Matrix
// Language: C++
#include <vector>
#include <algorithm>
using namespace std;
class Solution {
public:
vector<int> kWeakestRows(vector<vector<int>>& mat, int k)
{
// Step 1: Store the number of soldiers along with their row indices.
vector<pair<int, int>> rowStrength;
// Iterate over each row in the matrix.
for (int i = 0; i < mat.size(); i++)
{
// Count the number of soldiers (1s) in the current row.
int soldiers = count(mat[i].begin(), mat[i].end(), 1);
// Store the soldier count along with the row index.
rowStrength.push_back({soldiers, i});
}
// Step 2: Sort the rows based on the number of soldiers.
// If two rows have the same number of soldiers, sort by row index.
sort(rowStrength.begin(), rowStrength.end());
// Step 3: Extract the indices of the k weakest rows.
vector<int> result;
// Select the first k weakest rows from the sorted list.
for (int i = 0; i < k; i++)
{
// Store the row index in the result vector.
result.push_back(rowStrength[i].second);
}
// Step 4: Return the result vector containing the k weakest row indices.
return result;
}
};Java
// Approach: Sorting-Based Solution to Find the K Weakest Rows in a Matrix
// Language: Java
import java.util.*;
class Solution {
public int[] kWeakestRows(int[][] mat, int k)
{
// Step 1: Store the number of soldiers along with their row indices.
List<int[]> rowStrength = new ArrayList<>();
// Iterate over each row in the matrix.
for (int i = 0; i < mat.length; i++)
{
// Count the number of soldiers (1s) in the current row.
int soldiers = countSoldiers(mat[i]);
// Store the soldier count along with the row index.
rowStrength.add(new int[]{soldiers, i});
}
// Step 2: Sort the rows based on the number of soldiers.
// If two rows have the same number of soldiers, sort by row index.
Collections.sort(rowStrength, (a, b) ->
{
// Compare by number of soldiers first.
if (a[0] != b[0])
{
return Integer.compare(a[0], b[0]);
}
// If soldiers are equal, compare by row index.
return Integer.compare(a[1], b[1]);
});
// Step 3: Extract the indices of the k weakest rows.
int[] result = new int[k];
// Select the first k weakest rows from the sorted list.
for (int i = 0; i < k; i++)
{
// Store the row index in the result array.
result[i] = rowStrength.get(i)[1];
}
// Step 4: Return the result array containing the k weakest row indices.
return result;
}
// Helper method to count the number of soldiers (1s) in a row.
private int countSoldiers(int[] row)
{
int count = 0;
// Iterate through the row and count 1s.
for (int num : row)
{
if (num == 1)
{
count++;
}
else
{
// Stop counting when the first 0 is encountered.
break;
}
}
// Return the total count of soldiers in the row.
return count;
}
}
Python
# Approach: Sorting-Based Solution to Find the K Weakest Rows in a Matrix
# Language: Python
class Solution(object):
def kWeakestRows(self, mat, k):
"""
:type mat: List[List[int]]
:type k: int
:rtype: List[int]
"""
# Step 1: Store the number of soldiers along with their row indices.
row_strength = []
# Iterate over each row in the matrix.
for i in range(len(mat)):
# Count the number of soldiers (1s) in the current row.
soldiers = self.count_soldiers(mat[i])
# Store the soldier count along with the row index.
row_strength.append((soldiers, i))
# Step 2: Sort the rows based on the number of soldiers.
# If two rows have the same number of soldiers, sort by row index.
row_strength.sort()
# Step 3: Extract the indices of the k weakest rows.
result = []
# Select the first k weakest rows from the sorted list.
for i in range(k):
# Store the row index in the result list.
result.append(row_strength[i][1])
# Step 4: Return the result list containing the k weakest row indices.
return result
# Helper method to count the number of soldiers (1s) in a row.
def count_soldiers(self, row):
count = 0
# Iterate through the row and count 1s.
for num in row:
if num == 1:
count += 1
else:
# Stop counting when the first 0 is encountered.
break
# Return the total count of soldiers in the row.
return countJavascript
/**
* Approach: Sorting-Based Solution to Find the K Weakest Rows in a Matrix
* Language: JavaScript
*
* @param {number[][]} mat - The input matrix containing soldiers (1s) and civilians (0s).
* @param {number} k - The number of weakest rows to return.
* @return {number[]} - The indices of the k weakest rows.
*/
var kWeakestRows = function(mat, k) {
// Step 1: Store the number of soldiers along with their row indices.
let rowStrength = [];
// Iterate over each row in the matrix.
for (let i = 0; i < mat.length; i++) {
// Count the number of soldiers (1s) in the current row.
let soldiers = countSoldiers(mat[i]);
// Store the soldier count along with the row index.
rowStrength.push([soldiers, i]);
}
// Step 2: Sort the rows based on the number of soldiers.
// If two rows have the same number of soldiers, sort by row index.
rowStrength.sort((a, b) => {
// Compare by number of soldiers first.
if (a[0] !== b[0]) {
return a[0] - b[0];
}
// If soldiers are equal, compare by row index.
return a[1] - b[1];
});
// Step 3: Extract the indices of the k weakest rows.
let result = [];
// Select the first k weakest rows from the sorted list.
for (let i = 0; i < k; i++) {
// Store the row index in the result array.
result.push(rowStrength[i][1]);
}
// Step 4: Return the result array containing the k weakest row indices.
return result;
};
/**
* Helper function to count the number of soldiers (1s) in a row.
*
* @param {number[]} row - The row from the matrix.
* @return {number} - The count of soldiers in the row.
*/
function countSoldiers(row) {
let count = 0;
// Iterate through the row and count 1s.
for (let num of row) {
if (num === 1) {
count++;
} else {
// Stop counting when the first 0 is encountered.
break;
}
}
// Return the total count of soldiers in the row.
return count;
}Time Complexity Analysis: O(m log m + m n)
Let m be the number of rows and n be the number of columns in the matrix.
Counting Soldiers:
- For each row, we count the number of soldiers. In the worst case, this takes O(n) time per row.
- Since there are m rows, this step takes O(m n) time in total.
Sorting the Rows:
- We sort the m rows based on the soldier count and row index. Sorting takes O(m log m) time.
Extracting the k Weakest Rows:
- After sorting, we extract the first k row indices, which takes O(k) time. Since k ≤ m, this simplifies to O(m).
Final Time Complexity:
- The dominant factors are O(m n) for counting soldiers and O(m log m) for sorting.
- Thus, the overall time complexity is O(m log m + m n).
Final Time Complexity: O(m log m + m n).
Space Complexity Analysis: O(m)
Auxiliary Space Complexity: O(m)
- We use an extra array to store the number of soldiers along with row indices. This array has m elements, so it requires O(m) space.
- Sorting is done in place, but since we are storing row information separately, the additional space remains O(m).
Result Storage:
- We store the indices of the k weakest rows in a separate list, which takes O(k) space. Since k ≤ m, this simplifies to O(m) in the worst case.
Total Auxiliary Space Complexity: O(m)
- The input matrix mat is already given, and we do not modify it. We only use extra space for the row strength list and result storage.
Final Space Complexity: O(m).
Instead of sorting all rows, we use a min-heap to efficiently track the weakest rows based on soldier count. By inserting each row into the heap and extracting the k weakest, we achieve an optimized solution in O(m log m) time, improving performance over full sorting.
Better Approach
Intuition
To optimize the process of finding the k weakest rows, a Min-Heap (Priority Queue) can be used instead of sorting. The idea is to store each row along with its soldier count in a heap, where the weakest rows (those with fewer soldiers) naturally rise to the top. Since a heap maintains order efficiently, inserting each row into the heap ensures that extracting the k weakest rows is done in an optimal manner.
As we iterate through the matrix, we compute the number of soldiers in each row and push it into the heap along with the row index. Since a Min-Heap always keeps the smallest element at the top, the weakest rows are easily accessible. Once all rows are processed, we extract the top k elements from the heap, which represent the k weakest rows. This approach avoids the overhead of full sorting, leading to better efficiency when compared to direct sorting.
Approach
Step 1: Create a Min-Heap
- A min-heap (priority queue) is used to efficiently store and retrieve the k weakest rows.
- Each entry in the heap consists of a pair: the number of soldiers in a row and the row index.
- The heap ensures that the row with the fewest soldiers is always at the top.
Step 2: Insert All Rows into the Heap
- Iterate through each row in the matrix.
- Count the number of soldiers (1s) in the row, since they always appear before civilians (0s).
- Insert a pair (soldier count, row index) into the min-heap.
Step 3: Extract the k Weakest Rows
- The heap automatically maintains order based on soldier count.
- Extract the first k elements from the heap, as they represent the k weakest rows.
- Store the row indices in the result list.
Step 4: Return the Result
- The result list now contains the indices of the k weakest rows in order.
- Return this list as the final output.
Dry Run
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]
Step 1: Create a Min-Heap
- Initialize a min-heap to store row indices sorted by their soldier count.
Step 2: Insert All Rows into the Heap
- Pick 3 → Heap = [3]
- Pick 2 → Heap = [2, 3]
Step 3: Process the Remaining Elements
Count the number of soldiers in each row and insert them into the heap:
- Row 0 → Soldiers = 2 → Heap = [(2,0)]
- Row 1 → Soldiers = 4 → Heap = [(2,0), (4,1)]
- Row 2 → Soldiers = 1 → Heap = [(1,2), (4,1), (2,0)]
- Row 3 → Soldiers = 2 → Heap = [(1,2), (2,3), (2,0), (4,1)]
- Row 4 → Soldiers = 5 → Heap = [(1,2), (2,3), (2,0), (4,1), (5,4)]
Step 3: Extract the k Weakest Rows
Since the heap maintains the weakest rows at the top, extract the first k elements:
- Extract (1,2) → Result = [2]
- Extract (2,0) → Result = [2,0]
- Extract (2,3) → Result = [2,0,3]
Step 4: Return the Result
- The weakest k rows are [2,0,3], so the final output is:
Final Output: [2,0,3]
Code for All Languages
C++
// Approach: Heap-Based Solution to Find the K Weakest Rows
// Language: C++
#include <vector>
#include <queue>
using namespace std;
class Solution {
public:
vector<int> kWeakestRows(vector<vector<int>>& mat, int k)
{
// Step 1: Create a min-heap (priority queue).
// The heap stores pairs (number of soldiers, row index).
priority_queue<pair<int, int>, vector<pair<int, int>>, greater<>> minHeap;
// Step 2: Compute the number of soldiers in each row and insert into the heap.
for (int i = 0; i < mat.size(); i++)
{
// Count the number of soldiers (1s) in the row.
int soldiers = count(mat[i].begin(), mat[i].end(), 1);
// Push the (soldiers count, row index) pair into the min-heap.
minHeap.push({soldiers, i});
}
// Step 3: Extract the k weakest rows.
vector<int> result;
while (k-- > 0)
{
// The weakest row is at the top of the min-heap.
result.push_back(minHeap.top().second);
// Remove the weakest row from the heap.
minHeap.pop();
}
// Step 4: Return the result containing the k weakest row indices.
return result;
}
};Java
// Approach: Heap-Based Solution to Find the K Weakest Rows
// Language: Java
import java.util.PriorityQueue;
class Solution {
public int[] kWeakestRows(int[][] mat, int k)
{
// Step 1: Create a min-heap (priority queue).
// The heap stores pairs (number of soldiers, row index).
PriorityQueue<int[]> minHeap = new PriorityQueue<>(
(a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]
);
// Step 2: Compute the number of soldiers in each row and insert into the heap.
for (int i = 0; i < mat.length; i++)
{
// Count the number of soldiers (1s) in the row.
int soldiers = 0;
for (int num : mat[i])
{
if (num == 1) soldiers++;
else break; // Since 1s always appear before 0s, we can stop early.
}
// Push the (soldiers count, row index) pair into the min-heap.
minHeap.offer(new int[]{soldiers, i});
}
// Step 3: Extract the k weakest rows.
int[] result = new int[k];
for (int i = 0; i < k; i++)
{
// The weakest row is at the top of the min-heap.
result[i] = minHeap.poll()[1];
}
// Step 4: Return the result containing the k weakest row indices.
return result;
}
}Python
# Approach: Heap-Based Solution to Find the K Weakest Rows
# Language: Python
import heapq
class Solution(object):
def kWeakestRows(self, mat, k):
"""
:type mat: List[List[int]]
:type k: int
:rtype: List[int]
"""
# Step 1: Create a min-heap (priority queue).
# The heap stores tuples (number of soldiers, row index).
min_heap = []
# Step 2: Compute the number of soldiers in each row and insert into the heap.
for i, row in enumerate(mat):
# Count the number of soldiers (1s) in the row.
soldiers = sum(row) # Since 1s always come before 0s, sum counts all 1s.
# Push the (soldiers count, row index) tuple into the min-heap.
heapq.heappush(min_heap, (soldiers, i))
# Step 3: Extract the k weakest rows.
result = []
for _ in range(k):
# The weakest row is at the top of the min-heap.
result.append(heapq.heappop(min_heap)[1])
# Step 4: Return the result containing the k weakest row indices.
return resultJavascript
// Approach: Heap-Based Solution to Find the K Weakest Rows
// Language: JavaScript
/**
* @param {number[][]} mat
* @param {number} k
* @return {number[]}
*/
var kWeakestRows = function(mat, k) {
// Step 1: Create a min-heap (priority queue).
// The heap stores pairs (number of soldiers, row index).
let minHeap = [];
// Step 2: Compute the number of soldiers in each row and insert into the heap.
for (let i = 0; i < mat.length; i++) {
// Count the number of soldiers (1s) in the row.
let soldiers = mat[i].reduce((sum, num) => sum + num, 0);
// Push the (soldiers count, row index) pair into the min-heap.
minHeap.push([soldiers, i]);
}
// Step 3: Sort the heap based on the number of soldiers and row index.
minHeap.sort((a, b) => a[0] === b[0] ? a[1] - b[1] : a[0] - b[0]);
// Step 4: Extract the k weakest rows.
let result = [];
for (let i = 0; i < k; i++) {
// The weakest row is at the beginning of the sorted heap.
result.push(minHeap[i][1]);
}
// Step 5: Return the result containing the k weakest row indices.
return result;
};Time Complexity Analysis: O(m * n) + O(m log m)
Let m be the number of rows and n be the number of columns in the matrix.
Counting Soldiers in Each Row - O(m * n)
- To determine the strength of each row, we iterate through all m rows and count the number of soldiers (1s) in each row. Since counting takes O(n) time per row, the total time complexity for this step is O(m * n).
Building the Min-Heap - O(m log m)
- After computing soldier counts, we insert m rows into a min-heap (priority queue). Each insertion operation takes O(log m) time, leading to an overall complexity of O(m log m).
Extracting k Weakest Rows - O(k log m)
- We extract k elements from the heap, and each extraction takes O(log m) time. This contributes O(k log m) time. However, since k ≤ m, it is bounded by O(m log m) in the worst case.
Final Complexity:O(m * n) + O(m log m)
The O(m * n) term comes from counting soldiers, and O(m log m) arises from sorting the rows using a min-heap. Thus, the overall time complexity is O(m * n) + O(m log m).
Space Complexity Analysis: O(m)
Auxiliary Space Complexity: O(m)
- We store the soldier count and row index for each of the m rows, requiring O(m) additional space.
Heap Storage: O(m)
- The min-heap (priority queue) stores at most m elements, each containing a pair of values (soldier count and row index). Since each insertion requires O(1) space per row, the total space required for the heap is O(m).
Input Matrix Storage: O(1) (excluding input)
- The given matrix mat is already provided as input, so it does not contribute to extra space usage. We do not create additional copies of the matrix, ensuring that the space complexity remains independent of n.
Final Space Complexity: O(m)
- The dominant factor in space usage is the heap and auxiliary storage for row strengths, both of which contribute O(m) space. Hence, the total space complexity remains O(m).
Instead of sorting all rows, we use a min-heap to efficiently track the weakest rows based on soldier count. By inserting each row into the heap and extracting the k weakest, we achieve an optimized solution in O(m log m) time, improving performance over full sorting.
Optimal Approach
Intuition
To optimize the process of finding the k weakest rows, a Min-Heap (Priority Queue) can be used instead of sorting. The idea is to store each row along with its soldier count in a heap, where the weakest rows (those with fewer soldiers) naturally rise to the top. Since a heap maintains order efficiently, inserting each row into the heap ensures that extracting the k weakest rows is done in an optimal manner.
As we iterate through the matrix, we compute the number of soldiers in each row and push it into the heap along with the row index. Since a Min-Heap always keeps the smallest element at the top, the weakest rows are easily accessible. Once all rows are processed, we extract the top k elements from the heap, which represent the k weakest rows. This approach avoids the overhead of full sorting, leading to better efficiency when compared to direct sorting.
Approach
Step 1: Create a Min-Heap
- A min-heap (priority queue) is used to efficiently store and retrieve the k weakest rows.
- Each entry in the heap consists of a pair: the number of soldiers in a row and the row index.
- The heap ensures that the row with the fewest soldiers is always at the top.
Step 2: Insert All Rows into the Heap
- Iterate through each row in the matrix.
- Count the number of soldiers (1s) in the row, since they always appear before civilians (0s).
- Insert a pair (soldier count, row index) into the min-heap.
Step 3: Extract the k Weakest Rows
- The heap automatically maintains order based on soldier count.
- Extract the first k elements from the heap, as they represent the k weakest rows.
- Store the row indices in the result list.
Step 4: Return the Result
- The result list now contains the indices of the k weakest rows in order.
- Return this list as the final output.
Dry Run
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]
Step 1: Create a Min-Heap
- Initialize a min-heap to store row indices sorted by their soldier count.
Step 2: Insert All Rows into the Heap
- Pick 3 → Heap = [3]
- Pick 2 → Heap = [2, 3]
Step 3: Process the Remaining Elements
Count the number of soldiers in each row and insert them into the heap:
- Row 0 → Soldiers = 2 → Heap = [(2,0)]
- Row 1 → Soldiers = 4 → Heap = [(2,0), (4,1)]
- Row 2 → Soldiers = 1 → Heap = [(1,2), (4,1), (2,0)]
- Row 3 → Soldiers = 2 → Heap = [(1,2), (2,3), (2,0), (4,1)]
- Row 4 → Soldiers = 5 → Heap = [(1,2), (2,3), (2,0), (4,1), (5,4)]
Step 3: Extract the k Weakest Rows
Since the heap maintains the weakest rows at the top, extract the first k elements:
- Extract (1,2) → Result = [2]
- Extract (2,0) → Result = [2,0]
- Extract (2,3) → Result = [2,0,3]
Step 4: Return the Result
- The weakest k rows are [2,0,3], so the final output is:
Final Output: [2,0,3]
Code for All Languages
C++
// Approach: Binary Search + Sorting to Find the K Weakest Rows
// Language: C++
#include <vector>
#include <algorithm>
using namespace std;
class Solution {
public:
// Function to count the number of soldiers in a row using binary search
int countSoldiers(vector<int>& row)
{
return upper_bound(row.begin(), row.end(), 1, greater<int>()) - row.begin();
}
vector<int> kWeakestRows(vector<vector<int>>& mat, int k)
{
// Step 1: Compute the number of soldiers in each row
vector<pair<int, int>> rowStrength;
for (int i = 0; i < mat.size(); i++)
{
// Count the number of soldiers in the row using binary search
int soldiers = countSoldiers(mat[i]);
// Store (soldiers count, row index) in the vector
rowStrength.push_back({soldiers, i});
}
// Step 2: Sort the rows by the number of soldiers, breaking ties by row index
sort(rowStrength.begin(), rowStrength.end());
// Step 3: Extract the indices of the k weakest rows
vector<int> result;
for (int i = 0; i < k; i++)
{
result.push_back(rowStrength[i].second);
}
// Step 4: Return the result containing the k weakest row indices
return result;
}
};Java
// Approach: Binary Search + Sorting to Find the K Weakest Rows
// Language: Java
import java.util.*;
class Solution {
// Function to count the number of soldiers in a row using binary search
private int countSoldiers(int[] row) {
int left = 0, right = row.length;
// Perform binary search to find the number of soldiers
while (left < right) {
int mid = left + (right - left) / 2;
// If mid is 1, move right to find the first zero
if (row[mid] == 1) {
left = mid + 1;
}
// Otherwise, move left
else {
right = mid;
}
}
// Return the number of soldiers in the row
return left;
}
public int[] kWeakestRows(int[][] mat, int k) {
// Step 1: Create a list to store (soldier count, row index) pairs
List<int[]> rowStrength = new ArrayList<>();
// Step 2: Compute the number of soldiers in each row
for (int i = 0; i < mat.length; i++) {
// Count soldiers using binary search
int soldiers = countSoldiers(mat[i]);
// Store the (soldiers, index) pair
rowStrength.add(new int[]{soldiers, i});
}
// Step 3: Sort the rows based on soldier count
// If two rows have the same soldier count, sort by row index
rowStrength.sort((a, b) -> (a[0] == b[0]) ? a[1] - b[1] : a[0] - b[0]);
// Step 4: Extract the indices of the k weakest rows
int[] result = new int[k];
for (int i = 0; i < k; i++) {
// Store the index of the weakest row in the result array
result[i] = rowStrength.get(i)[1];
}
// Step 5: Return the result containing the k weakest row indices
return result;
}
}Python
# Approach: Heap-Based Solution to Find the K Weakest Rows
# Language: Python
import heapq
class Solution(object):
def kWeakestRows(self, mat, k):
"""
:type mat: List[List[int]]
:type k: int
:rtype: List[int]
"""
# Step 1: Create a min-heap (priority queue).
# The heap stores tuples (number of soldiers, row index).
min_heap = []
# Step 2: Compute the number of soldiers in each row and insert into the heap.
for i, row in enumerate(mat):
# Count the number of soldiers (1s) in the row.
soldiers = sum(row) # Since 1s always come before 0s, sum counts all 1s.
# Push the (soldiers count, row index) tuple into the min-heap.
heapq.heappush(min_heap, (soldiers, i))
# Step 3: Extract the k weakest rows.
result = []
for _ in range(k):
# The weakest row is at the top of the min-heap.
result.append(heapq.heappop(min_heap)[1])
# Step 4: Return the result containing the k weakest row indices.
return resultJavascript
// Approach: Heap-Based Solution to Find the K Weakest Rows
// Language: JavaScript
/**
* @param {number[][]} mat
* @param {number} k
* @return {number[]}
*/
var kWeakestRows = function(mat, k) {
// Step 1: Create a min-heap (priority queue).
// The heap stores pairs (number of soldiers, row index).
let minHeap = [];
// Step 2: Compute the number of soldiers in each row and insert into the heap.
for (let i = 0; i < mat.length; i++) {
// Count the number of soldiers (1s) in the row.
let soldiers = mat[i].reduce((sum, num) => sum + num, 0);
// Push the (soldiers count, row index) pair into the min-heap.
minHeap.push([soldiers, i]);
}
// Step 3: Sort the heap based on the number of soldiers and row index.
minHeap.sort((a, b) => a[0] === b[0] ? a[1] - b[1] : a[0] - b[0]);
// Step 4: Extract the k weakest rows.
let result = [];
for (let i = 0; i < k; i++) {
// The weakest row is at the beginning of the sorted heap.
result.push(minHeap[i][1]);
}
// Step 5: Return the result containing the k weakest row indices.
return result;
};Time Complexity Analysis: O(m * n) + O(m log m)
Let m be the number of rows and n be the number of columns in the matrix.
Counting Soldiers in Each Row - O(m * n)
- To determine the strength of each row, we iterate through all m rows and count the number of soldiers (1s) in each row. Since counting takes O(n) time per row, the total time complexity for this step is O(m * n).
Building the Min-Heap - O(m log m)
- After computing soldier counts, we insert m rows into a min-heap (priority queue). Each insertion operation takes O(log m) time, leading to an overall complexity of O(m log m).
Extracting k Weakest Rows - O(k log m)
- We extract k elements from the heap, and each extraction takes O(log m) time. This contributes O(k log m) time. However, since k ≤ m, it is bounded by O(m log m) in the worst case.
Final Complexity:O(m * n) + O(m log m)
The O(m * n) term comes from counting soldiers, and O(m log m) arises from sorting the rows using a min-heap. Thus, the overall time complexity is O(m * n) + O(m log m).
Space Complexity Analysis: O(m)
Auxiliary Space Complexity: O(m)
- We store the soldier count and row index for each of the m rows, requiring O(m) additional space.
Heap Storage: O(m)
- The min-heap (priority queue) stores at most m elements, each containing a pair of values (soldier count and row index). Since each insertion requires O(1) space per row, the total space required for the heap is O(m).
Input Matrix Storage: O(1) (excluding input)
- The given matrix mat is already provided as input, so it does not contribute to extra space usage. We do not create additional copies of the matrix, ensuring that the space complexity remains independent of n.
Final Space Complexity: O(m)
- The dominant factor in space usage is the heap and auxiliary storage for row strengths, both of which contribute O(m) space. Hence, the total space complexity remains O(m).
Similar Problems
Now we have successfully tackled this problem, let’s try these similar problems.
Given an integer array nums and an integer k, return the kth largest element in the array.Note that it is the kth largest element in the sorted order, not the kth distinct element.
This problem requires continuously finding the median of a growing data stream. It uses two heaps (a max-heap for the left half and a min-heap for the right half) to efficiently balance and retrieve the median in O(log n) insertion and O(1) retrieval time, similar to how we maintain the k-th largest element in a min-heap.