Program to find how many ways we can climb stairs (maximum steps at most k times) in Python

Suppose we have a staircase with n steps and we also have another number k. Initially we are at stair 0, and we can climb up either 1, 2 or 3 steps at a time, but we can only climb 3 stairs at most k times. We need to find the number of ways we can climb the staircase.

So, if the input is like n = 5, k = 2, then the output will be 13, as there are different ways we can climb the stairs −

• [1, 1, 1, 1, 1]
• [2, 1, 1, 1]
• [1, 2, 1, 1]
• [1, 1, 2, 1]
• [1, 1, 1, 2]
• [1, 2, 2]
• [2, 1, 2]
• [2, 2, 1]
• [1, 1, 3]
• [1, 3, 1]
• [3, 1, 1]
• [2, 3]
• [3, 2]

Algorithm

To solve this problem, we use dynamic programming with a 2D memoization table. The approach follows these steps −

• Handle base cases: if n is 0 or 1, return 1
• Create a memo table where memo[j][i] represents ways to reach step i using at most j three-step moves
• Fill the memo table using recurrence relations based on whether we can use more three-step moves

Example

class Solution:
    def solve(self, n, k):
        # Base cases
        if n == 0:
            return 1
        if n == 1:
            return 1
        
        # Optimize k to avoid unnecessary computation
        k = min(k, n)
        
        # Create memoization table: memo[j][i] = ways to reach step i using at most j three-steps
        memo = [[0] * (n + 1) for _ in range(k + 1)]
        
        # Initialize base cases for all rows
        for r in range(k + 1):
            memo[r][0] = 1  # One way to stay at step 0
            memo[r][1] = 1  # One way to reach step 1
            if n >= 2:
                memo[r][2] = 2  # Two ways to reach step 2: [1,1] or [2]
        
        # Fill first row (no three-steps allowed)
        for i in range(3, n + 1):
            memo[0][i] = memo[0][i-1] + memo[0][i-2]
        
        # Fill remaining rows
        for j in range(1, k + 1):
            for i in range(3, n + 1):
                count = i // 3  # Maximum three-steps needed to reach step i
                if count <= j:
                    # We have enough three-step moves available
                    memo[j][i] = memo[j][i-1] + memo[j][i-2] + memo[j][i-3]
                else:
                    # We need to use one less three-step move from previous state
                    memo[j][i] = memo[j][i-1] + memo[j][i-2] + memo[j-1][i-3]
        
        return memo[k][n]

# Test the solution
ob = Solution()
result = ob.solve(n=5, k=2)
print(f"Number of ways to climb {5} stairs with at most {2} three-step moves: {result}")

Number of ways to climb 5 stairs with at most 2 three-step moves: 13

How It Works

The algorithm uses a 2D dynamic programming table where:

• memo[j][i] represents the number of ways to reach step i using at most j three-step moves
• For each position, we consider three possibilities: coming from 1 step back, 2 steps back, or 3 steps back
• The constraint on three-step moves is handled by tracking how many we've used and adjusting the recurrence accordingly

Time and Space Complexity

• Time Complexity: O(n × k) where n is the number of steps and k is the maximum three-step moves allowed
• Space Complexity: O(n × k) for the memoization table

Conclusion

This dynamic programming solution efficiently counts all valid ways to climb stairs with step size constraints. The key insight is using a 2D table to track both position and remaining three-step moves.