Minimum cost in a 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

Assume we have already choose elements from first $i$ rows. The only thing we care about is current sum (smaller is better), column we took in the last row, and number of changes we already did.

So we process rows, supporting state min_sum[number of changes][last column taken], and the answer will be min(min_sum[K]). Can you find transition function from $i$-th to $i+1$-th row?