299红色格子染色方案数计算

问题描述

小 R 有一排长度为 n 的格子，每个格子从左到右编号为 1 到 n。起初，部分格子已经被染成了红色，其他格子则没有颜色。红色格子的状态由一个长度为 n 的字符串 s 描述，其中 s[i] = 1 表示第 i 个格子是红色的，而 s[i] = 0 表示该格子没有颜色。

小 R 希望通过以下两种操作将所有格子都染成红色：

1. 如果第 i 个格子是红色的，且 i + 1 ≤ n，则可以将第 i + 1 个没有颜色的格子染成红色。
2. 如果第 i 个格子是红色的，且 i - 1 ≥ 1，则可以将第 i - 1 个没有颜色的格子染成红色。

请你帮小 R 计算出，存在多少种不同的染色顺序可以使所有格子最终都被染成红色，并输出答案对 10^9 + 7 取模后的结果。

题解

1.子数组内部的染色方案数
首先考虑一个长度为k的全部未染色的格子，可以有多少种染色方案。

• 假如这段连续的格子处于两个已染色的格子之间，那么这段格子的染色方案数为2^(k-1)。
• 假如这段连续的格子只有一端被染色，那么显然只有一种染色方案。

2. 合并所有的方案数

• 假设有x段连续的未染色格子，那么可以发现，这x段格子的染色方案数是独立的，可以相乘。问题相当于把这x段格子放入长度为m的数组中，m为未染色的格子总数。
• 例如 有长度为x1、x2、x3、x4的未染色格子，一共有m个未染色格子，放入其中的方案数为C(m, x1) * C(m-x1, x2) * C(m-x1-x2, x3) * C(m-x1-x2-x3, x4)。

mod = int(1e9 + 7)
mx = int(1e5 + 10)

# 预处理组合数
fac = [0] * mx
fac[0] = 1
for i in range(1, mx):
    fac[i] = fac[i - 1] * i % mod

inv_fac = [0] * mx
inv_fac[mx - 1] = pow(fac[mx - 1], -1, mod)
for i in range(mx - 1, 0, -1):
    inv_fac[i - 1] = inv_fac[i] * i % mod

def comb(n, m):
    if n < m:
        return 0
    return fac[n] * inv_fac[m] % mod * inv_fac[n - m] % mod

def solution(n, s):
    a = [i for i in range(n) if s[i] == '1']
    # 未染色的格子数
    cnt = n - len(a)
    # 处理首尾
    res = comb(cnt, a[0]) * comb(cnt - a[0], n - a[-1] - 1) % mod
    cnt -= (a[0] + n - a[-1] - 1)
    for i in range(1, len(a)):
        k = a[i] - a[i - 1] - 1
        if k > 0:
            # 内部的方案数和放置的方案数相乘
            res *= pow(2, k - 1, mod) * comb(cnt, k) % mod
            cnt -= k
    return res % mod


if __name__ == '__main__':
    print(solution(n = 5,s = "00101") == 3)
    print(solution(n = 6,s = "100001") == 8)
    print(solution(n = 7,s = "0001000") == 20)

复杂度分析

• 时间复杂度：O(n)，可以预处理pow
• 空间复杂度：O(n)