LC 2954. 统计感冒序列的数目

https://leetcode.cn/problems/count-the-number-of-infection-sequences/description/

这题初看还有点难,但是细看是个比较简单的组合问题,可能唯一麻烦一点的就是需要求解一下乘法逆元,不过好在我有模板。

sick之间分为两种情况处理:

接着就是把多个区间的人开始做排列组合,假设有4个区间分别是a,b,c,d的话,那么排列组合的可能性是

\[A(a+b+c+d) / A(a) / A(b) / A(c) / A(d)\] 所有人之间假设可以任意排列,但是a,b,c,d内部不能任意排列,其中 \(A(x) = x!\)

#!/usr/bin/env python
# coding:utf-8
# Copyright (C) dirlt

from typing import List


# 费马小定理, 但是这里必须确保MOD是质数
# b^MOD % MOD = b
# b^(MOD-1) % MOD = 1
# b^(MOD-2) % MOD = (b^-1) % MOD
def pow_mod(a, b, MOD):
    res = 1
    while b:
        if b & 0x1:
            res = (res * a) % MOD
        a = (a * a) % MOD
        b = b >> 1
    return res


def div_mod(b, MOD):
    return pow_mod(b, MOD - 2, MOD)


def fac_mod(n, MOD):
    res = 1
    for i in range(1, n + 1):
        res = (res * i) % MOD
    return res


class Solution:
    def numberOfSequence(self, n: int, sick: List[int]) -> int:
        MOD = 10 ** 9 + 7

        gap = []
        middle = []
        for i in range(1, len(sick)):
            g = sick[i] - sick[i - 1] - 1
            if g > 0:
                middle.append(g)
                gap.append(g)

        if sick[0] != 0:
            gap.append(sick[0])
        if sick[-1] != n - 1:
            gap.append(n - 1 - sick[-1])

        # print(gap, middle)

        ans = 1
        for x in middle:
            ans *= pow_mod(2, x - 1, MOD)
            ans %= MOD

        ans *= fac_mod(sum(gap), MOD)
        ans %= MOD

        for g in gap:
            r = fac_mod(g, MOD)
            ans *= div_mod(r, MOD)
            ans %= MOD

        return ans