找出所有稳定的二进制数组I
题目描述1
题目讲解
为什么要用 DP?
因为你每一步都要决定:
- 放 0 还是放 1 ?
- 放了之后会不会超过 limit ?
- 还剩多少 0 和 1 ?
这就是典型的“状态依赖前一步”的问题 → 用 DP。
我们只关心:
- 当前用了多少个 0
- 当前用了多少个 1
- 最后放的是 0 还是 1
所以我们定义:
dp0[i][j] = 用了 i 个 0、j 个 1,并且最后放的是 0 的方案数
dp1[i][j] = 用了 i 个 0、j 个 1,并且最后放的是 1 的方案数
假设我们要算: dp0[i][j](最后放 0)
那前一步必须是:
- 最后放 1(否则连续 0 会变长)
- 并且我们可以连续放 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