LeetCode-day34-3130. 找出所有稳定的二进制数组 II

题目描述

给你 3 个正整数 zero ，one 和 limit 。

一个 二进制数组 arr 如果满足以下条件，那么我们称它是 稳定的 ：

• 0 在 arr 中出现次数 恰好 为 zero 。
• 1 在 arr 中出现次数 恰好 为 one 。
• arr 中每个长度超过 limit 的 子数组 都 同时 包含 0 和 1 。

请你返回 稳定 二进制数组的 总 数目。

由于答案可能很大，将它对 10⁹ + 7 取余 后返回。

示例

示例1：

输入：zero = 1, one = 1, limit = 2
输出：2
解释：
两个稳定的二进制数组为 [1,0] 和 [0,1] ，两个数组都有一个 0 和一个 1 ，且没有子数组长度大于 2 。

示例2：

输入：zero = 1, one = 2, limit = 1
输出：1
解释：
唯一稳定的二进制数组是 [1,0,1] 。
二进制数组 [1,1,0] 和 [0,1,1] 都有长度为 2 且元素全都相同的子数组，所以它们不稳定。

示例3：

输入：zero = 3, one = 3, limit = 2
输出：14
解释：
所有稳定的二进制数组包括 [0,0,1,0,1,1] ，[0,0,1,1,0,1] ，[0,1,0,0,1,1] ，[0,1,0,1,0,1] ，[0,1,0,1,1,0] ，[0,1,1,0,0,1] ，[0,1,1,0,1,0] ，[1,0,0,1,0,1] ，[1,0,0,1,1,0] ，[1,0,1,0,0,1] ，[1,0,1,0,1,0] ，[1,0,1,1,0,0] ，[1,1,0,0,1,0] 和 [1,1,0,1,0,0] 。

思路

采用递推解题
状态转移方程为：

f[i][j][0] = f[i - 1][j][0] + f[i - 1][j][1] - f[i - limit - 1][j][1]
f[i][j][1] = f[i][j - 1][0] + f[i][j - 1][1] - f[i][j - limit - 1][0]

代码

class Solution:
    def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
        MOD = 1_000_000_007
        f = [[[0, 0] for _ in range(one + 1)] for _ in range(zero + 1)]
        for i in range(1, min(limit, zero) + 1):
            f[i][0][0] = 1
        for j in range(1, min(limit, one) + 1):
            f[0][j][1] = 1
        for i in range(1, zero + 1):
            for j in range(1, one + 1):
                f[i][j][0] = (f[i - 1][j][0] + f[i - 1][j][1] - (f[i - limit - 1][j][1] if i > limit else 0)) % MOD
                f[i][j][1] = (f[i][j - 1][0] + f[i][j - 1][1] - (f[i][j - limit - 1][0] if j > limit else 0)) % MOD
        return sum(f[-1][-1]) % MOD