BinarySearch 1003 Drop a Ball Down the Stairs

Problem statement

https://binarysearch.com/problems/Drop-a-Ball-Down-the-Stairs/

Solution

We can use dp here with states (i, step).

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, H, blacklist):
        b_set, M = set(blacklist), 10**9 + 7
        x = [[1, 2, 4], [1, 3, 4]]

        @lru_cache(None)
        def dp(i, step):
            if i in b_set or i <= -1: return 0
            if i == 0: return 1
            return sum([dp(i - x, 1 - step) for x in x[step]])

        return dp(H, 0) % M