The robot has three choices: stay in place, move one step to the left, or move one step to the right.
How many ways are there to keep the robot at the origin (index 0) after steps steps?
1 <= steps <= 5001 <= arrLen <= 10^6
The dependency is the sum of pos-1, pos, and pos+1 from the previous step.
Naive DFS: TLE
java
class Solution {
public int numWays(int steps, int arrLen) {
return dfs(0, steps, arrLen);
}
int dfs(int pos, int steps, int arrLen) {
if (steps == 0) {
return pos == 0 ? 1 : 0;
}
int mod = (int) 1e9 + 7;
int res = 0;
res = (res + dfs(pos, steps - 1, arrLen)) % mod;
if (pos > 0) res = (res + dfs(pos - 1, steps - 1, arrLen)) % mod;
if (pos + 1 < arrLen) res = (res + dfs(pos + 1, steps - 1, arrLen)) % mod;
return res;
}
}Optimization is necessary: to return to index 0, the robot cannot wander too far. Note that arrLen is at the million level!
We should return dp[steps][0] because we need the number of ways to reach position 0.
Iterate starting from step = 1.
DP
java
class Solution {
public int numWays(int steps, int arrLen) {
int n = Math.min(steps / 2 + 1, arrLen);
int mod = (int) 1e9 + 7;
long[][] dp = new long[steps + 1][n];
dp[0][0] = 1;
for (int i = 1; i < steps + 1; i++) {
for (int j = 0; j < n; j++) {
dp[i][j] = (dp[i - 1][j] + (j + 1 < n ? dp[i - 1][j + 1] : 0) + (j > 0 ? dp[i - 1][j - 1] : 0)) % mod;
}
}
return (int) dp[steps][0];
}
}Space compression
java
class Solution {
public int numWays(int steps, int arrLen) {
int n = Math.min(steps / 2 + 1, arrLen);
int mod = (int) 1e9 + 7;
long[][] dp = new long[2][n];
dp[0][0] = 1;
for (int i = 1; i < steps + 1; i++) {
for (int j = 0; j < n; j++) {
dp[i & 1][j] = (dp[i - 1 & 1][j] + (j + 1 < n ? dp[i - 1 & 1][j + 1] : 0) + (j > 0 ? dp[i - 1 & 1][j - 1] : 0)) % mod;
}
}
return (int) dp[steps & 1][0];
}
}