LeetCode 2954. Count the Number of Infection Sequences

December 3, 2023

周賽374。直接來一個純數學題，還以為是dp，直接暴死。
雖說已經連續7場周賽沒有AK，但是分數竟然是上漲的，可見最近是真的難。

題目

輸入整數n，還有遞增排序過的整數陣列sick。

有n個小孩排隊站在地點0~n-1上，而sick陣列代表了生病小孩的位置。如果地點i的小孩生病了，他可以傳染給i-1或是i+1的小孩。每一秒鐘只能有一個小孩被感染。

過了一段時間後，所有的小孩都會被感染，而感染序列指的是由那些被依序感染的小孩的位置。
求共有多少可能的感染序列。答案可能很大，先模10^9+7後回傳。

注意：感染序列不包含最初就被感染的小孩。

解法

至少有一個染感的人在隊伍中。以X代表感染者、o代表健康，隊伍會有兩種情況：

o…X 一群健康人只有單邊接觸感染者
Xo..oX 兩個感染者夾著一群健康人

第一種情況只會出現在第一人或最後一人健康的時候出現。感染順序也很單純，只能單方向感染。
第二種情況占多數，可感染可選右邊或左邊。例如XoooX選左邊就變成XXooX、選右邊變成XooXX。
其實就是變成兩個規模縮減的子問題，基本情形就是只有一人被夾著，只有一種選法。有n個人被夾，就有2^(n-1)種選法。

那麼問題來了，如果有兩區以上的健康人怎麼辦，這兩種如何交錯處理？
例如：XooXoX裡面有兩區，分別以A表示第一區的人，B表示第二區的人，以AAB表示。
我們先想想第一區的人要在何時感染。現在要在3個格子內填入2個A，總共有C(3,2)種組合：

空格子[_,_,_]填2個A
[A,A,_]
[A,_,A]
[_,A,A]
共三種組合

那麼對於每種組合，又可以分別帶入他的感染順序。以a1, a2代表第一區的兩個人代入：

[A,A,_]這種組合可以有[a1,a2,_]和[a2,a1,_]兩種排列
[A,_,A]這種組合可以有[a1,_,a2]和[a2,_,a2]兩種排列
[_,A,A]這種組合可以有[_,a1,a2]和[_,a2,a1]兩種排列

所以3個格子中填兩個A，共有C(3,2) * 2^(2-1)種 = 3 * 2種排列。
剩下1個格子填入一個B，共有C(1,1) * 2^(1-1)種 = 1 * 1種排列。
答案就是6 * 1種排列。

至於第一種情況，同樣是從空格中選填，只不過只有一種排列方式而已。

先預處理所有會用到的階乘，以及用於除法的模反元素(乘法逆元)。

時間複雜度O(m log n)，其中m為sick長度。
空間複雜度O(1)，預處理空間不計入。

MOD=10**9+7
MX=10**5+5
f=[1]*MX
for i in range(2,MX):
    f[i]=f[i-1]*i
    f[i]%=MOD

rev=[1]*MX
for i in range(2,MX):
    rev[i]=pow(f[i],-1,MOD)
    
def C(n,k):
    return f[n]*rev[k]*rev[n-k]

class Solution:
    def numberOfSequence(self, n: int, sick: List[int]) -> int:
        ans=1
        left=sick[0]
        right=n-1-sick[-1]
        tot=left+right
        for a,b in pairwise(sick):
            tot+=b-a-1
        
        for a,b in pairwise(sick):
            x=b-a-1
            if x>0:
                ans*=C(tot,x)*pow(2,x-1,MOD)
                tot-=x
        
        for x in [left,right]:
            if x>0:
                ans*=C(tot,x)
                tot-=x
        
        return ans%MOD

mocowcow

題目

解法