Program to check number of ways we can move k times and return back to first place in python

Suppose we are at position 0 of an n-length array, and on each step, we can move right one place, left one place (not exceeding bounds), or stay at the same position. We need to find how many unique walks of exactly k steps return us back to index 0. If the answer is very large, return it modulo 10^9 + 7.

So, if the input is like n = 7, k = 4, then the output will be 9. The possible action sequences are ?

• [Right, Right, Left, Left]
• [Right, Left, Right, Left]
• [Stay, Stay, Stay, Stay]
• [Right, Left, Stay, Stay]
• [Stay, Stay, Right, Left]
• [Right, Stay, Stay, Left]
• [Right, Stay, Left, Stay]
• [Stay, Right, Left, Stay]
• [Stay, Right, Stay, Left]

Algorithm Approach

We'll use dynamic programming with memoization to solve this problem efficiently. The approach follows these steps ?

• Define a recursive function dp(position, steps_remaining)
• Base case: if no steps remain, return 1 if at position 0, else 0
• For each position, try all three moves: stay, move left, move right
• Sum up all valid paths and return the result modulo 10^9 + 7

Implementation

class Solution:
    def solve(self, length, k):
        MOD = 10**9 + 7
        memo = {}
        
        def dp(position, steps_remaining):
            # Base case: no steps left
            if steps_remaining == 0:
                return 1 if position == 0 else 0
            
            # Memoization check
            if (position, steps_remaining) in memo:
                return memo[(position, steps_remaining)]
            
            count = 0
            
            # Stay at current position
            count += dp(position, steps_remaining - 1)
            
            # Move right (if within bounds)
            if position + 1 < length:
                count += dp(position + 1, steps_remaining - 1)
            
            # Move left (if within bounds)
            if position - 1 >= 0:
                count += dp(position - 1, steps_remaining - 1)
            
            count %= MOD
            memo[(position, steps_remaining)] = count
            return count
        
        return dp(0, k)

# Test the solution
solution = Solution()
n = 7
k = 4
result = solution.solve(n, k)
print(f"Number of ways to return to position 0 in {k} steps: {result}")

Number of ways to return to position 0 in 4 steps: 9

How It Works

The algorithm uses memoization to avoid recalculating the same subproblems. For each position and remaining steps, we explore three possibilities ?

• Stay: Remain at current position, use one step
• Move Right: Go to position + 1 (if valid), use one step
• Move Left: Go to position - 1 (if valid), use one step

Alternative Approach Using Bottom-Up DP

def count_ways_iterative(n, k):
    MOD = 10**9 + 7
    
    # dp[step][pos] = number of ways to reach position pos in step steps
    dp = [[0] * n for _ in range(k + 1)]
    dp[0][0] = 1  # Base case: 1 way to be at position 0 with 0 steps
    
    for step in range(1, k + 1):
        for pos in range(n):
            # Stay at current position
            dp[step][pos] += dp[step - 1][pos]
            
            # Come from left position
            if pos > 0:
                dp[step][pos] += dp[step - 1][pos - 1]
            
            # Come from right position
            if pos < n - 1:
                dp[step][pos] += dp[step - 1][pos + 1]
            
            dp[step][pos] %= MOD
    
    return dp[k][0]

# Test the iterative solution
n = 7
k = 4
result = count_ways_iterative(n, k)
print(f"Iterative approach result: {result}")

Iterative approach result: 9

Comparison

Conclusion

This problem demonstrates classic dynamic programming where we build solutions from smaller subproblems. Both recursive and iterative approaches work efficiently with O(n × k) time complexity, making them suitable for reasonable input sizes.