Count ways to reach end from start stone and come back, without taking using the same stone twice with at most K jumps at each step

Question

Richard is a penguin who regularly commutes between Canada and the United States.

Specifically, Canada is at location 1 , and the United States is at location N. In locations 2 through N, there are ice blocks that Richard can jump on. Richard can jump from location i to location j if and only if abs(i-j)<=k.

Richard does not believe in wasting time, so when he is commuting in one direction, he will always move in that direction. This means that when commuting from Canada to the United States, locations are monotonically increasing in number, and when commuting from the United States to Canada, locations are monotonically decreasing in number.

Richard's commute is complicated during the fall and spring when the ice is thawing between Canada and the United States. If Richard jumps on an ice block, that ice block will melt enough that it will not be usable on the return journey.

Count the number of distinct ways that Richard can do a round-trip commute from Canada to the United States and back. Two ways are distinct if, in some direction, Richard uses an ice block travelling in that direction in one way but not in the other way.

2<=k<=N<=5000

sample input: 5 3 sample output 12

how should i approach this problem?

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)

here is my code for calculating the way there, but what about the way back?your text

Answer 1

Imagine coloring rocks 2 through N-1. Color them red if they're used on the way from Canada to the US, and color them blue if they're used on the way back. Leave them gray if they are not used at all. Your question is how many ways you can color the rocks such that there are no gaps of >= k non-red stones or non-blue stones.

Make a matrix DP[r][b] that contains the number of ways to color the stones such that the highest red stone is r, the highest blue stone is b, and there are no illegal red gaps < r or blue gaps < b.

If you fill out the matrix in order of increasing r+b, max(r,b), etc., then each element can be calculated in O(k) time from previously filled out elements. Then your final answer can be computed in O(k²) time given the filled out matrix.

Since the matrix has size (N-2)², this takes O(kN²) time.