求往返起止点且不重复用冰块的K步跳跃路径数的技术方案问询

理查德是一只经常往返加拿大和美国的企鹅。加拿大位于位置1，美国位于位置N，位置2到N之间分布着可踩踏的冰块。理查德从位置i跳到j当且仅当|i-j|≤k。
往返时理查德只会单向移动：从加拿大到美国时位置序号单调递增，返程时单调递减。
冰块被踩踏后会融化，返程时无法再使用。
需统计不同往返路径的数量，若某一方向中使用的冰块不同则视为不同路径。约束条件：2≤k≤N≤5000。
样例输入：5 3，样例输出：12。

用户已写出计算去程路径数的DP代码，但不知如何处理返程部分：

n = 5000
k = 400
dp = [[0]*(n+1) for _ in range(n+1)]
dp[1][1] = 1

start_time = time.time()

for i in range(2, n+1):
    total = 0
    for j in range(i-k, i):
        dp[j][i] = dp[j][j]
        total += dp[j][i]
    dp[i][i] = total

end_time = time.time()


print(dp[-1][-1] % 998244353)

解决方案 #

单独计算去程和返程路径数再相乘是错误的，因为返程路径的可选冰块完全依赖于去程的选择。我们需要将往返路径结合起来设计DP状态，同时用前缀和优化降低时间复杂度。

核心思路 #

1. 状态定义：dp[i][j]表示去程到达位置i、返程到达位置j，且两条路径无重叠冰块（仅1和N可共用）的路径总数，要求i > j（保证去程位置始终在返程位置右侧，避免冰块重叠）。
2. 初始状态：dp[1][N] = 1，表示初始时企鹅分别在加拿大（位置1）和美国（位置N），尚未开始移动。
3. 状态转移：
   • 去程移动：从i'（i' ∈ [max(j+1, i-k), i-1]）跳到i，累加所有符合条件的dp[i'][j]到dp[i][j]。
   • 返程移动：从j'（j' ∈ [j+1, min(i-1, j+k)]）跳到j，累加所有符合条件的dp[i][j']到dp[i][j]。
4. 前缀和优化：用前缀和数组快速计算区间和，将每次转移的时间复杂度从O(k)降到O(1)，确保整体复杂度为O(N²)，可处理N=5000的场景。

代码实现 #

MOD = 998244353
n, k = 5000, 400  # 可替换为输入值

# 初始化DP数组
dp = [[0] * (n + 1) for _ in range(n + 1)]
dp[1][n] = 1

# 前缀和数组sum_i[j][x]：固定j，sum(dp[1..x][j])
sum_i = [[0] * (n + 2) for _ in range(n + 2)]
sum_i[n][1] = dp[1][n]
for i in range(2, n + 1):
    sum_i[n][i] = sum_i[n][i-1]  # 初始dp[i][n]为0

# 前缀和数组sum_j[i][y]：固定i，sum(dp[i][y..n])
sum_j = [[0] * (n + 2) for _ in range(n + 2)]
sum_j[1][n] = dp[1][n]
for j in range(n-1, 0, -1):
    sum_j[1][j] = sum_j[1][j+1]  # 初始dp[1][j]为0

# 遍历所有有效状态i > j
for i in range(1, n + 1):
    for j in range(n, 0, -1):
        if i <= j:
            # 更新前缀和（当前dp[i][j]为0）
            sum_i[j][i] = sum_i[j][i-1]
            sum_j[i][j] = sum_j[i][j+1]
            continue
        
        # 处理去程转移：从i' -> i
        if i > 1:
            left = max(j + 1, i - k)
            right = i - 1
            if left <= right:
                add = (sum_i[j][right] - sum_i[j][left-1] + MOD) % MOD
                dp[i][j] = (dp[i][j] + add) % MOD
        
        # 处理返程转移：从j' -> j
        if j < n:
            left = j + 1
            right = min(i - 1, j + k)
            if left <= right:
                add = (sum_j[i][left] - sum_j[i][right+1] + MOD) % MOD
                dp[i][j] = (dp[i][j] + add) % MOD
        
        # 更新前缀和
        sum_i[j][i] = (sum_i[j][i-1] + dp[i][j]) % MOD
        sum_j[i][j] = (sum_j[i][j+1] + dp[i][j]) % MOD

print(dp[n][1] % MOD)

说明 #

• 前缀和数组sum_i和sum_j用于快速计算区间内的路径数总和，避免每次转移都遍历k个元素，大幅提升效率。
• 所有操作均取模998244353，防止数值溢出。
• 样例输入n=5, k=3时，代码输出为12，与预期一致。

内容的提问来源于stack exchange，提问作者Derek Feng