题目描述1

题目讲解

为什么要用 DP?

因为你每一步都要决定:

  1. 放 0 还是放 1 ?
  2. 放了之后会不会超过 limit ?
  3. 还剩多少 0 和 1 ?

这就是典型的“状态依赖前一步”的问题 → 用 DP。

我们只关心:

  1. 当前用了多少个 0
  2. 当前用了多少个 1
  3. 最后放的是 0 还是 1

所以我们定义:

dp0[i][j] = 用了 i 个 0、j 个 1,并且最后放的是 0 的方案数
dp1[i][j] = 用了 i 个 0、j 个 1,并且最后放的是 1 的方案数

假设我们要算: dp0[i][j](最后放 0)

那前一步必须是:

  1. 最后放 1(否则连续 0 会变长)
  2. 并且我们可以连续放 1 个、2 个、…、limit 个 0

即:

dp0[i][j] = dp1[i-1][j] + dp1[i-2][j] + ... + dp1[i-limit][j]
+ 解释 dp0[i][j] = dp1[i-1][j] + dp1[i-2][j] + ... + dp1[i-limit][j]

假设:i = 5, j = 3, limit = 2 , 你要算 dp0[5][3](最后是 0)

最后一段 0 的长度可以是:1 个 0, 2 个 0, 不能超过 limit=2.

最后一段是 1 个 0:
... 1 0
前面用了 4 个 0
→ dp1[4][3]
最后一段是 2 个 0:
... 1 0 0
前面用了 3 个 0
→ dp1[3][3]

所以:dp0[5][3] = dp1[4][3] + dp1[3][3]

同理:

dp1[i][j] = dp0[i][j-1] + dp0[i][j-2] + ... + dp0[i][j-limit]

为什么要前缀和?

每次都要加 limit 项, 太慢. 前缀和可以把: dp1[i-1][j] + dp1[i-2][j] + ... + dp1[i-limit][j] 变成 pre1[i-1][j] - pre1[i-limit-1][j] .

+ 解释 dp1[i-1][j] + dp1[i-2][j] + ... + dp1[i-limit][j] 变成 pre1[i-1][j] - pre1[i-limit-1][j]

假设你有一排数字:

a[0], a[1], a[2], a[3], a[4], ...

前缀和 pre 定义为:

pre[k] = a[0] + a[1] + ... + a[k]

pre[r] = a[0] + a[1] + ... + a[l-1] + a[l] + ... + a[r]
pre[l-1] = a[0] + a[1] + ... + a[l-1]
a[l] + a[l+1] + ... + a[r] = pre[r] - pre[l-1]

即得到前缀和公式

区间和 = pre1[右端点] - pre1[左端点 - 1]

代码

python不使用前缀和版
python使用前缀和版
class Solution:
    def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
        MOD = 10**9 + 7

        # dp0[i][j]:用了 i 个 0、j 个 1,最后一段是 0
        # dp1[i][j]:用了 i 个 0、j 个 1,最后一段是 1
        dp0 = [[0] * (one + 1) for _ in range(zero + 1)]
        dp1 = [[0] * (one + 1) for _ in range(zero + 1)]

        # 初始化:只放 0 或只放 1 的情况
        for i in range(1, min(zero, limit) + 1):
            dp0[i][0] = 1
        for j in range(1, min(one, limit) + 1):
            dp1[0][j] = 1

        for i in range(zero + 1):
            for j in range(one + 1):
                if i == 0 and j == 0:
                    continue

                # 计算 dp0[i][j]:最后一段是 0,长度为 k(1..limit)
                if i > 0:
                    val = 0
                    for k in range(1, min(limit, i) + 1):
                        # 前面用了 i-k 个 0、j 个 1,且最后一段必须是 1
                        val += dp1[i - k][j]
                    dp0[i][j] = (dp0[i][j] + val) % MOD

                # 计算 dp1[i][j]:最后一段是 1,长度为 k(1..limit)
                if j > 0:
                    val = 0
                    for k in range(1, min(limit, j) + 1):
                        # 前面用了 i 个 0、j-k 个 1,且最后一段必须是 0
                        val += dp0[i][j - k]
                    dp1[i][j] = (dp1[i][j] + val) % MOD

        return (dp0[zero][one] + dp1[zero][one]) % MOD
class Solution:
    def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
        MOD = 10**9 + 7

        dp0 = [[0] * (one + 1) for _ in range(zero + 1)]
        dp1 = [[0] * (one + 1) for _ in range(zero + 1)]

        # pre1[i][j] = dp1[0][j] + ... + dp1[i][j]   (沿着 i 累加)
        # pre0[i][j] = dp0[i][0] + ... + dp0[i][j]   (沿着 j 累加)
        pre1 = [[0] * (one + 1) for _ in range(zero + 1)]
        pre0 = [[0] * (one + 1) for _ in range(zero + 1)]

        # 边界初始化:只有 0 或只有 1 的情况
        for i in range(1, min(zero, limit) + 1):
            dp0[i][0] = 1
        for j in range(1, min(one, limit) + 1):
            dp1[0][j] = 1

        # 顺便把前缀和也初始化好
        for i in range(zero + 1):
            for j in range(one + 1):
                # 更新 pre1(沿 i 累加)
                pre1[i][j] = ((pre1[i-1][j] if i > 0 else 0) + dp1[i][j]) % MOD
                # 更新 pre0(沿 j 累加)
                pre0[i][j] = ((pre0[i][j-1] if j > 0 else 0) + dp0[i][j]) % MOD

        for i in range(zero + 1):
            for j in range(one + 1):
                if i == 0 and j == 0:
                    continue

                # 只在 i>0 且 j>0 时,用转移公式更新
                # 边界(i>0,j=0 或 i=0,j>0)保留初始化结果
                if i > 0 and j > 0:
                    # dp0[i][j] = sum_{k=1..limit} dp1[i-k][j]
                    L = max(0, i - limit)
                    R = i - 1
                    dp0[i][j] = (pre1[R][j] - (pre1[L-1][j] if L > 0 else 0)) % MOD

                    # dp1[i][j] = sum_{k=1..limit} dp0[i][j-k]
                    L = max(0, j - limit)
                    R = j - 1
                    dp1[i][j] = (pre0[i][R] - (pre0[i][L-1] if L > 0 else 0)) % MOD

                # 更新当前格子的前缀和
                pre1[i][j] = ((pre1[i-1][j] if i > 0 else 0) + dp1[i][j]) % MOD
                pre0[i][j] = ((pre0[i][j-1] if j > 0 else 0) + dp0[i][j]) % MOD

        return (dp0[zero][one] + dp1[zero][one]) % MOD

链接