golang 解法, 执行用时: 7 ms, 内存消耗: 6.8 MB, 提交时间: 2024-05-01 10:52:33

func numberOfStableArrays(zero, one, limit int) (ans int) {
    const mod = 1_000_000_007
    f := make([][][2]int, zero+1)
    for i := range f {
        f[i] = make([][2]int, one+1)
    }
    for i := 1; i <= min(limit, zero); i++ {
        f[i][0][0] = 1
    }
    for j := 1; j <= min(limit, one); j++ {
        f[0][j][1] = 1
    }
    for i := 1; i <= zero; i++ {
        for j := 1; j <= one; j++ {
            f[i][j][0] = (f[i-1][j][0] + f[i-1][j][1]) % mod
            if i > limit {
                // + mod 保证答案非负
                f[i][j][0] = (f[i][j][0] - f[i-limit-1][j][1] + mod) % mod
            }
            f[i][j][1] = (f[i][j-1][0] + f[i][j-1][1]) % mod
            if j > limit {
                f[i][j][1] = (f[i][j][1] - f[i][j-limit-1][0] + mod) % mod
            }
        }
    }
    return (f[zero][one][0] + f[zero][one][1]) % mod
}

cpp 解法, 执行用时: 11 ms, 内存消耗: 13.6 MB, 提交时间: 2024-05-01 10:52:20

class Solution {
public:
    int numberOfStableArrays(int zero, int one, int limit) {
        const int MOD = 1'000'000'007;
        vector<vector<array<int, 2>>> f(zero + 1, vector<array<int, 2>>(one + 1));
        for (int i = 1; i <= min(limit, zero); i++) {
            f[i][0][0] = 1;
        }
        for (int j = 1; j <= min(limit, one); j++) {
            f[0][j][1] = 1;
        }
        for (int i = 1; i <= zero; i++) {
            for (int j = 1; j <= one; j++) {
                // + MOD 保证答案非负
                f[i][j][0] = ((long long) f[i - 1][j][0] + f[i - 1][j][1] + (i > limit ? MOD - f[i - limit - 1][j][1] : 0)) % MOD;
                f[i][j][1] = ((long long) f[i][j - 1][0] + f[i][j - 1][1] + (j > limit ? MOD - f[i][j - limit - 1][0] : 0)) % MOD;
            }
        }
        return (f[zero][one][0] + f[zero][one][1]) % MOD;
    }
};

java 解法, 执行用时: 18 ms, 内存消耗: 43.8 MB, 提交时间: 2024-05-01 10:51:26

class Solution {
    public int numberOfStableArrays(int zero, int one, int limit) {
        final int MOD = 1_000_000_007;
        int[][][] f = new int[zero + 1][one + 1][2];
        for (int i = 1; i <= Math.min(limit, zero); i++) {
            f[i][0][0] = 1;
        }
        for (int j = 1; j <= Math.min(limit, one); j++) {
            f[0][j][1] = 1;
        }
        for (int i = 1; i <= zero; i++) {
            for (int j = 1; j <= one; j++) {
                // + MOD 保证答案非负
                f[i][j][0] = (int) (((long) f[i - 1][j][0] + f[i - 1][j][1] + (i > limit ? MOD - f[i - limit - 1][j][1] : 0)) % MOD);
                f[i][j][1] = (int) (((long) f[i][j - 1][0] + f[i][j - 1][1] + (j > limit ? MOD - f[i][j - limit - 1][0] : 0)) % MOD);
            }
        }
        return (f[zero][one][0] + f[zero][one][1]) % MOD;
    }
}

python3 解法, 执行用时: 398 ms, 内存消耗: 28.9 MB, 提交时间: 2024-05-01 10:50:38

class Solution:
    def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
        MOD = 1_000_000_007
        @cache  # 缓存装饰器，避免重复计算 dfs 的结果（记忆化）
        
        # dfs(i,j,k) 表示用 i 个 0 和 j 个 1 构造稳定数组的方案数，
        # 其中第 i+j 个位置要填 k，其中 k 为 0 或 1。
        def dfs(i: int, j: int, k: int) -> int:
            if i == 0:
                return 1 if k == 1 and j <= limit else 0
            if j == 0:
                return 1 if k == 0 and i <= limit else 0
            if k == 0:
                return (dfs(i - 1, j, 0) + dfs(i - 1, j, 1) - (dfs(i - limit - 1, j, 1) if i > limit else 0)) % MOD
            else:  # else 可以去掉，这里仅仅是为了代码对齐
                return (dfs(i, j - 1, 0) + dfs(i, j - 1, 1) - (dfs(i, j - limit - 1, 0) if j > limit else 0)) % MOD
        ans = (dfs(zero, one, 0) + dfs(zero, one, 1)) % MOD
        dfs.cache_clear()  # 防止爆内存
        return ans
        
    
    # 递推
    def numberOfStableArrays2(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