3406. Find All Possible Stable Binary Arrays I

class Solution:
    def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
        MOD = 10**9 + 7
        
        # dp[z][o][last] = number of stable arrays with z zeros, o ones, ending with last (0 or 1)
        # We use inclusion-exclusion to handle the limit constraint
        
        # Let's think differently:
        # f(z, o, last) = ways to place z zeros and o ones ending with 'last'
        # where no consecutive run exceeds limit
        
        # For each position, we can extend a run of the same type (if < limit)
        # or start a new run of the opposite type
        
        # Let dp[z][o][0] = ways ending with 0
        # Let dp[z][o][1] = ways ending with 1
        
        # Transition:
        # dp[z][o][0] = sum over len=1 to min(limit, z) of dp[z-len][o][1] (place len zeros after some ones)
        # dp[z][o][1] = sum over len=1 to min(limit, o) of dp[z][o-len][0] (place len ones after some zeros)
        
        # Base case: dp[0][0][*] is not valid, we need at least one element
        # dp[k][0][0] = 1 for k in 1..min(limit, zero) (just zeros)
        # dp[0][k][1] = 1 for k in 1..min(limit, one) (just ones)
        
        dp = [[[0, 0] for _ in range(one + 1)] for _ in range(zero + 1)]
        
        # Base cases: arrays with only zeros or only ones
        for z in range(1, min(limit, zero) + 1):
            dp[z][0][0] = 1
        for o in range(1, min(limit, one) + 1):
            dp[0][o][1] = 1
        
        # Fill the DP table
        for z in range(zero + 1):
            for o in range(one + 1):
                if z == 0 and o == 0:
                    continue
                    
                # Ending with 0: we place 1 to min(limit, z) zeros after something ending with 1
                for length in range(1, min(limit, z) + 1):
                    if z - length >= 0 and (z - length > 0 or o > 0):
                        if z - length == 0 and o > 0:
                            dp[z][o][0] = (dp[z][o][0] + dp[0][o][1]) % MOD
                        elif z - length > 0:
                            dp[z][o][0] = (dp[z][o][0] + dp[z - length][o][1]) % MOD
                
                # Ending with 1: we place 1 to min(limit, o) ones after something ending with 0
                for length in range(1, min(limit, o) + 1):
                    if o - length >= 0 and (o - length > 0 or z > 0):
                        if o - length == 0 and z > 0:
                            dp[z][o][1] = (dp[z][o][1] + dp[z][0][0]) % MOD
                        elif o - length > 0:
                            dp[z][o][1] = (dp[z][o][1] + dp[z][o - length][0]) % MOD
        
        return (dp[zero][one][0] + dp[zero][one][1]) % MOD