Find Valid Matrix Given Row and Column Sums

medium array greedy matrix

Problem

Given rowSum and colSum of nonneg ints, return any nonneg matrix whose row sums and col sums match. Such a matrix always exists.

InputrowSum = [3,8], colSum = [4,7]

Output[[3,0],[1,7]]

Place min(3,4)=3 at (0,0); row 0 finished. Then min(8,1)=1 at (1,0). Then 7 at (1,1).

rowSum (comma-separated)

colSum

def restore_matrix(rowSum, colSum):
    m, n = len(rowSum), len(colSum)
    M = [[0]*n for _ in range(m)]
    i = j = 0
    while i < m and j < n:
        v = min(rowSum[i], colSum[j])
        M[i][j] = v
        rowSum[i] -= v
        colSum[j] -= v
        if rowSum[i] == 0:
            i += 1
        else:
            j += 1
    return M

function restoreMatrix(rowSum, colSum) {
  const m = rowSum.length, n = colSum.length;
  const M = Array.from({length: m}, () => Array(n).fill(0));
  let i = 0, j = 0;
  while (i < m && j < n) {
    const v = Math.min(rowSum[i], colSum[j]);
    M[i][j] = v;
    rowSum[i] -= v; colSum[j] -= v;
    if (rowSum[i] === 0) i++; else j++;
  }
  return M;
}

class Solution {
    public int[][] restoreMatrix(int[] rowSum, int[] colSum) {
        int m = rowSum.length, n = colSum.length;
        int[][] M = new int[m][n];
        int i = 0, j = 0;
        while (i < m && j < n) {
            int v = Math.min(rowSum[i], colSum[j]);
            M[i][j] = v;
            rowSum[i] -= v; colSum[j] -= v;
            if (rowSum[i] == 0) i++; else j++;
        }
        return M;
    }
}

vector<vector<int>> restoreMatrix(vector<int>& rowSum, vector<int>& colSum) {
    int m = rowSum.size(), n = colSum.size();
    vector<vector<int>> M(m, vector<int>(n, 0));
    int i = 0, j = 0;
    while (i < m && j < n) {
        int v = min(rowSum[i], colSum[j]);
        M[i][j] = v;
        rowSum[i] -= v; colSum[j] -= v;
        if (rowSum[i] == 0) i++; else j++;
    }
    return M;
}

Explanation

We need to build any matrix whose row totals and column totals match the given rowSum and colSum. The neat insight is that a valid answer always exists, so we can be greedy and fill cells one at a time without ever getting stuck.

We keep a pointer i for the current row and j for the current column. At cell (i, j) we place min(rowSum[i], colSum[j]) — the largest amount that neither overshoots the remaining row budget nor the remaining column budget.

After placing that value v, we subtract it from both rowSum[i] and colSum[j]. One of them must now hit zero: if the row is used up we move down (i += 1), otherwise the column is done so we move right (j += 1). Every other cell we never touch stays at its default 0.

Example: rowSum = [3, 8], colSum = [4, 7]. Place min(3,4)=3 at (0,0); row 0 is finished, move down. Place min(8,1)=1 at (1,0); column 0 finished, move right. Place min(7,7)=7 at (1,1). Result: [[3,0],[1,7]].

Because each step finishes either a row or a column, we only take about m + n steps total.

Time: O(m + n) Space: O(m·n)