题目地址:
https://leetcode.com/problems/cherry-pickup-ii/
给定一个二维矩阵 A A A,元素非负。有两个机器人,分别从第 0 0 0行的首尾两个格子开始向下走,每次可以走到左下、正下或者右下一格内,不能出界。问两个机器人走到最后一行的时候,路径上的数字总和最大能达到多少。如果两个机器人走到了相同的格子,则数字只能取一次。
思路是动态规划。设
f
[
i
]
[
j
]
[
k
]
f[i][j][k]
f[i][j][k]表示当两个机器人走到了第
i
i
i行,第一个机器人走到第
j
j
j列,第二个走到第
k
k
k列时,最大的路径数字和。则有递推关系
f
[
i
]
[
j
]
[
k
]
=
{
max
j
′
∈
{
j
−
1
,
j
,
j
+
1
}
,
k
′
∈
{
k
−
1
,
k
,
k
+
1
}
{
f
[
i
−
1
]
[
j
′
]
[
k
′
]
}
+
A
[
i
]
[
j
]
,
j
=
k
max
j
′
∈
{
j
−
1
,
j
,
j
+
1
}
,
k
′
∈
{
k
−
1
,
k
,
k
+
1
}
{
f
[
i
−
1
]
[
j
′
]
[
k
′
]
}
+
A
[
i
]
[
j
]
+
A
[
i
]
[
k
]
,
j
≠
k
f[i][j][k] =\begin{cases} \max_{j'\in\{j-1,j,j+1\}, k'\in\{k-1,k,k+1\}}\{f[i-1][j'][k']\}+ A[i][j],j=k\\ \max_{j'\in\{j-1,j,j+1\}, k'\in\{k-1,k,k+1\}}\{f[i-1][j'][k']\}+ A[i][j]+A[i][k],j\ne k \end{cases}
f[i][j][k]={maxj′∈{j−1,j,j+1},k′∈{k−1,k,k+1}{f[i−1][j′][k′]}+A[i][j],j=kmaxj′∈{j−1,j,j+1},k′∈{k−1,k,k+1}{f[i−1][j′][k′]}+A[i][j]+A[i][k],j=k逐层递推即可。
代码实现方面要注意两点:
1、由于一开始机器人的位置是给定的,所以向下走的时候有的状态是走不到的。例如
A
A
A是长度为
3
3
3宽为
4
4
4的矩阵,那么显然
f
[
1
]
[
1
]
[
3
]
f[1][1][3]
f[1][1][3]不能由
f
[
0
]
[
1
]
[
3
]
+
A
[
1
]
[
3
]
f[0][1][3]+A[1][3]
f[0][1][3]+A[1][3]去更新,因为一开始左边的机器人就不在
A
[
0
]
[
1
]
A[0][1]
A[0][1],这会造成更新出错误的答案。所以一开始我们要以
−
1
-1
−1记录某个状态不可达,然后每次只用可达的状态更新别的状态即可。
2、显然第一个机器人到达坐标
i
i
i的行的时候,其列坐标最大也就是
i
i
i,同理第二个机器人到达坐标
i
i
i的行的时候,其列坐标最小也就是
n
−
i
−
1
n-i-1
n−i−1。我们可以据此对代码进行优化,把那些不可达的状态排除掉。
代码如下:
import java.util.Arrays;
public class Solution {
public int cherryPickup(int[][] grid) {
int m = grid.length, n = grid[0].length;
int[][][] dp = new int[m][n][n];
// 每个位置初始化为-1
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
Arrays.fill(dp[i][j], -1);
}
}
// 第0行的状态初始化一下
dp[0][0][n - 1] = n != 1 ? grid[0][0] + grid[0][n - 1] : grid[0][0];
// i代表两个机器人走到的行坐标
for (int i = 1; i < m; i++) {
// j代表第一个机器人走到的列坐标,其范围可以优化
for (int j = 0; j <= Math.min(i, n - 1); j++) {
// k代表第一个机器人走到的列坐标,其范围也可以优化
for (int k = Math.max(0, n - i - 1); k <= n - 1; k++) {
// 开始枚举上一层的可达状态
for (int d1 = -1; d1 <= 1; d1++) {
for (int d2 = -1; d2 <= 1; d2++) {
int idx1 = j + d1, idx2 = k + d2;
// 如果没出界,并且不等于-1,说明可达,更新dp数组
if (0 <= idx1 && idx1 < n && 0 <= idx2 && idx2 < n) {
if (dp[i - 1][idx1][idx2] != -1) {
if (j == k) {
dp[i][j][k] = Math.max(dp[i][j][k], dp[i - 1][idx1][idx2] + grid[i][j]);
} else {
dp[i][j][k] = Math.max(dp[i][j][k], dp[i - 1][idx1][idx2] + grid[i][j] + grid[i][k]);
}
}
}
}
}
}
}
}
// 答案就是dp[m - 1]这个正方形矩阵里的最大值
int res = Integer.MIN_VALUE;
for (int i = 0; i < n; i++) {
for (int j = 0; j < n; j++) {
res = Math.max(res, dp[m - 1][i][j]);
}
}
return res;
}
}
时空复杂度 O ( m n 2 ) O(mn^2) O(mn2)。
C++:
class Solution {
public:
int cherryPickup(vector<vector<int>>& g) {
int m = g.size(), n = g[0].size();
int f[m][n][n];
memset(f, -1, sizeof f);
f[0][0][n - 1] = g[0][0];
if (n > 1) f[0][0][n - 1] += g[0][n - 1];
for (int k = 0; k < m - 1; k++)
for (int i = 0; i < n; i++)
for (int j = 0; j < n; j++)
if (~f[k][i][j])
for (int di = -1; di <= 1; di++)
for (int dj = -1; dj <= 1; dj++) {
int ni = i + di, nj = j + dj;
if (0 <= ni && ni < n && 0 <= nj && nj < n) {
int x = g[k + 1][ni];
if (ni != nj) x += g[k + 1][nj];
f[k + 1][ni][nj] = max(f[k + 1][ni][nj], f[k][i][j] + x);
}
}
int res = 0;
for (int i = 0; i < n; i++)
for (int j = 0; j < n; j++) res = max(res, f[m - 1][i][j]);
return res;
}
};
时空复杂度一样。