Finding minimum cost in 2D matrix

Question

In my last interview, I was asked a question for which optimal approach I am still not able to figure out.

Given a 2D matrix, with n rows and m columns, I need to select one item from each row and get the minimum sum.

But the tricky part was, while going down, I have two options:

1. choose the element in same column as prev one.
2. Choose element apart from the current column

So, the 2nd option, i.e changing column while going down can be done exactly K times. We can revisit the same column down the matrix. Change is defined only when we change the column between current and next row.

What I tried in front of the interviewer was something like Painting wall problem, where given 3 types of paints I need to find minimum cost so that no two adjacant walls are painted the same color. But there the m value was fixed to 3, and I can use simple cost[i][j] = min(cost[i-1][j-1], cost[i-1][j+1]); //min of the other two rows from the current one.

So problem as above, what should be the approach, it does feel like DP to me but I am really not sure if it gives result.

Answer 1

This can be solved using top down approach.

You start at a location on the first row and to traverse down from there. You have two options:

a) Continue on the same column - but you need to take into consideration that you have enough rows for K switches. So your condition should be remaining_rows >= K then you can take this option otherwise not.

b) Switch your column to any other column. Consider all columns.

Out of those two options take the minimum.

The recursive equation will be then:

getMin(r, c, K) = min( 
                   (R - (r+1) >= K ? getMin(r+1, c, K) + A[r][c] : A[r][c] ),
                   ( for all j != c, j < C: min(getMin(r+1, c, K-1)) )
                  )

where A is the input matrix and R is the number of rows and C is the number of columns in A.

and you take the minimum of getMin(0, c, K ) for all c = 0 to C

You can memoize the solution for each r, c, K.

For K = 0 , just traverse down the column and get the sum.

Code in Java is :

class Main {
    private static int[][][] dp;

    public static void main(String[] args) {

        int[][] A = { { 0, 1, 2 }, { 4, 3, 1 }, { 6, -2, 5 }, { 1, 4, -3 } };

        int K = 3;
        dp = new int[A.length][A[0].length][K + 1];
        for (int i = 0; i < A.length; i++) {
            for (int j = 0; j < A[0].length; j++) {
                for (int k = 1; k <= K; k++) {
                    dp[i][j][k] = Integer.MAX_VALUE;
                }
            }
        }
        int res = Integer.MAX_VALUE;
        for (int j = 0; j < A[0].length; j++)
            res = Math.min(res, getMin(A, 0, j, K));
        System.out.println(res);
    }

    private static int getMin(int[][] A, int r, int c, int K) {
        int R = A.length, C = A[0].length;
        if (r > R - 1 || c > C - 1)
            return Integer.MAX_VALUE;
        if (K == 0) {
            int sum = 0;
            for (int i = r; i < R; i++)
                sum += A[i][c];
            return sum;
        }
        if (dp[r][c][K] != Integer.MAX_VALUE)
            return dp[r][c][K];
        int min = Integer.MAX_VALUE;
        int rem_rows = R - (r + 1);
        for (int j = 0; j < C; j++) {
            if (j == c && rem_rows >= K) {
                int s = getMin(A, r + 1, j, K);
                if (s != Integer.MAX_VALUE) {
                    min = Math.min(min, s + A[r][c]);
                } else {
                    min = Math.min(min, A[r][c]);
                }

            } else if (j != c) {
                int m = getMin(A, r + 1, j, K - 1);
                if (m != Integer.MAX_VALUE) {
                    min = Math.min(min, m + A[r][c]);
                } else {
                    min = Math.min(min, A[r][c]);
                }
            }
        }
        dp[r][c][K] = min;
        return min;
    }
}

You can play with the live code here

Answer 2

This can be solved using dynamic programming: let dp[i][j][l] be the minimum sum we can get by choosing some numbers from the first i rows, the last one being in j th column, and having made exactly l column changes. Initially dp[0][j][0] = matrix[0][j]. For the updates we should the following formula:

If we go with the first option, let's do this:
dp[i + 1][j][l] = min(dp[i + 1][j][l], dp[i][j][l] + matrix[i + 1][j]).

If we go with the second option, we should do this for all new_col != j:
dp[i + 1][new_col][l + 1] = min(dp[i + 1][new_col][l + 1], dp[i][j][l] + matrix[i + 1][new_col])

The result will be the minimum value among dp[n - 1][j][k] for all j. Overall complexity is
O(n * m^2 * k) because for all states of the dynamic (O(n * m * k)) we do updates in O(m) time while iterating over new_col for the second type and O(1) for the first type.

We can do better than this. Let's find the smallest dp[i][j][l] for fixed i, l, iterating over all possible j. Suppose it's in column number min_j, and the value is x. For all j != min_j x will be used for updates of the second type in row i + 1.

The only column which is updated using a different value is min_j itself, because if we went from
dp[i][min_j][l] to dp[i + 1][min_j][l + 1], we wouldn't switch the column. Let's calculate
dp[i + 1][min_j][l + 1] by iterating over all j != min_j as we used to do before. Total number of operations is O(m).

The complexity will be O(n * m * k) because we are doing O(m) updates for all different pairs i, l and there's O(n * k) of them.