ABC156-D_Bouquetをpythonで解いた（フェルマーの小定理）

問題

• D - Bouquet

ポイント

• 組み合わせ、10⁹+7で割った余り
• -> フェルマーの小定理

MOD=10**9+7
def comb(n,k):
    if n<k:
        return 0
    if n<0 or k<0:
        return 0
    k=min(n-k,k)
    ans=1
    inv=[1]*(k+1)
    if k>=1:
        ans*=(n-k+1)%MOD
    for i in range(2,k+1):
        inv[i]=MOD-inv[MOD%i]*(MOD//i)%MOD
        ans=ans*(n-k+i)*inv[i]%MOD
    return ans

comb(4,2)
# -> 6

解法

n,a,b = map(int,input().split())

MOD=10**9+7
def comb(n,k):
    if n<k:
        return 0
    if n<0 or k<0:
        return 0
    k=min(n-k,k)
    ans=1
    inv=[1]*(k+1)
    if k>=1:
        ans*=(n-k+1)%MOD
    for i in range(2,k+1):
        inv[i]=MOD-inv[MOD%i]*(MOD//i)%MOD
        ans=ans*(n-k+i)*inv[i]%MOD
    return ans

nCa = comb(n,a)
nCb = comb(n,b)

ans = (pow(2,n,MOD) - 1) - nCa - nCb

if ans < 0:
    print((ans+MOD)%MOD)
else:
    print(ans%MOD)

感想

• これならnが10⁹でも計算できる。