Question

1269. 停在原地的方案数

English Version

题目描述

有一个长度为 arrLen 的数组，开始有一个指针在索引 0 处。

每一步操作中，你可以将指针向左或向右移动 1 步，或者停在原地（指针不能被移动到数组范围外）。

给你两个整数 steps 和 arrLen ，请你计算并返回：在恰好执行 steps 次操作以后，指针仍然指向索引 0 处的方案数。

由于答案可能会很大，请返回方案数 模 10^9 + 7 后的结果。

 

示例 1：

输入：steps = 3, arrLen = 2
输出：4
解释：3 步后，总共有 4 种不同的方法可以停在索引 0 处。
向右，向左，不动
不动，向右，向左
向右，不动，向左
不动，不动，不动

示例  2：

输入：steps = 2, arrLen = 4
输出：2
解释：2 步后，总共有 2 种不同的方法可以停在索引 0 处。
向右，向左
不动，不动

示例 3：

输入：steps = 4, arrLen = 2
输出：8

 

提示：

• 1 <= steps <= 500
• 1 <= arrLen <= 106

解法

方法一：记忆化搜索

我们观察题目的数据范围，可以发现 $steps$ 最大不超过 $500$，这意味着我们最远只能往右走 $500$ 步。

我们可以设计一个函数 $dfs(i, j)$，表示当前在位置 $i$，并且剩余步数为 $j$ 的方案数。那么答案就是 $dfs(0, steps)$。

函数 $dfs(i, j)$ 的执行过程如下：

1. 如果 $i \gt j$ 或者 $i \geq arrLen$ 或者 $i \lt 0$ 或者 $j \lt 0$，那么返回 $0$。
2. 如果 $i = 0$ 且 $j = 0$，那么此时指针已经停在原地，并且没有剩余步数，所以返回 $1$。
3. 否则，我们可以选择向左走一步，向右走一步，或者不动，所以返回 $dfs(i - 1, j - 1) + dfs(i + 1, j - 1) + dfs(i, j - 1)$。注意答案的取模操作。

过程中，我们可以使用记忆化搜索避免重复计算。

时间复杂度 $O(steps \times steps)$，空间复杂度 $O(steps \times steps)$。其中 $steps$ 是题目给定的步数。

class Solution:
  def numWays(self, steps: int, arrLen: int) -> int:
    @cache
    def dfs(i, j):
      if i > j or i >= arrLen or i < 0 or j < 0:
        return 0
      if i == 0 and j == 0:
        return 1
      ans = 0
      for k in range(-1, 2):
        ans += dfs(i + k, j - 1)
        ans %= mod
      return ans

    mod = 10**9 + 7
    return dfs(0, steps)

class Solution {
  private Integer[][] f;
  private int n;

  public int numWays(int steps, int arrLen) {
    f = new Integer[steps][steps + 1];
    n = arrLen;
    return dfs(0, steps);
  }

  private int dfs(int i, int j) {
    if (i > j || i >= n || i < 0 || j < 0) {
      return 0;
    }
    if (i == 0 && j == 0) {
      return 1;
    }
    if (f[i][j] != null) {
      return f[i][j];
    }
    int ans = 0;
    final int mod = (int) 1e9 + 7;
    for (int k = -1; k <= 1; ++k) {
      ans = (ans + dfs(i + k, j - 1)) % mod;
    }
    return f[i][j] = ans;
  }
}

class Solution {
public:
  int numWays(int steps, int arrLen) {
    int f[steps][steps + 1];
    memset(f, -1, sizeof f);
    const int mod = 1e9 + 7;
    function<int(int, int)> dfs = [&](int i, int j) -> int {
      if (i > j || i >= arrLen || i < 0 || j < 0) {
        return 0;
      }
      if (i == 0 && j == 0) {
        return 1;
      }
      if (f[i][j] != -1) {
        return f[i][j];
      }
      int ans = 0;
      for (int k = -1; k <= 1; ++k) {
        ans = (ans + dfs(i + k, j - 1)) % mod;
      }
      return f[i][j] = ans;
    };
    return dfs(0, steps);
  }
};

func numWays(steps int, arrLen int) int {
  const mod int = 1e9 + 7
  f := make([][]int, steps)
  for i := range f {
    f[i] = make([]int, steps+1)
    for j := range f[i] {
      f[i][j] = -1
    }
  }
  var dfs func(i, j int) int
  dfs = func(i, j int) (ans int) {
    if i > j || i >= arrLen || i < 0 || j < 0 {
      return 0
    }
    if i == 0 && j == 0 {
      return 1
    }
    if f[i][j] != -1 {
      return f[i][j]
    }
    for k := -1; k <= 1; k++ {
      ans += dfs(i+k, j-1)
      ans %= mod
    }
    f[i][j] = ans
    return
  }
  return dfs(0, steps)
}

function numWays(steps: number, arrLen: number): number {
  const f = Array.from({ length: steps }, () => Array(steps + 1).fill(-1));
  const mod = 10 ** 9 + 7;
  const dfs = (i: number, j: number) => {
    if (i > j || i >= arrLen || i < 0 || j < 0) {
      return 0;
    }
    if (i == 0 && j == 0) {
      return 1;
    }
    if (f[i][j] != -1) {
      return f[i][j];
    }
    let ans = 0;
    for (let k = -1; k <= 1; ++k) {
      ans = (ans + dfs(i + k, j - 1)) % mod;
    }
    return (f[i][j] = ans);
  };
  return dfs(0, steps);
}