AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC172 E - NEQ

問題原文

atcoder.jp

解法

まず全パターンは、 「 M 個の数から  N 個選んで並べる」を「 A B に対してやる」 ので、  ({}_M P_N)^{2} 通り。
ここから、ダメなパターンを引いていくことを考える。

ダメなパターンは「1つが一致しているパターン」「2つが一致しているパターン」「3つが。。。」の形になってるので、なんか包除原理が使えそうだなとなる。

考えるべきことは以下。
例えばここでは、  r 個一致しているパターンを考える。

  •  M 個の数から  r 個選ぶ
  •  N 個の置き場所から  r 個選んで並べる
  • 残り  M - r 個から、  N - r 個選んで並べる。これを  A,  B に対してやる。

すると以下のようなコードになる。↓

実装

N, M = map(int, input().split())
MOD = 10 ** 9 + 7

factorial = [1, 1]
inverse = [1, 1]
invere_base = [0, 1]
for i in range(2, N + M + 1):
    factorial.append((factorial[-1] * i) % MOD)
    invere_base.append((-invere_base[MOD % i] * (MOD // i)) % MOD)
    inverse.append((inverse[-1] * invere_base[-1]) % MOD)

def nCr(n, r):
    if not (0 <= r <= n):
        return 0
    return factorial[n] * inverse[r] * inverse[n - r] % MOD

def nPr(n, r):
    if not (0 <= r <= n):
        return 0
    return nCr(n, r) * factorial[r] % MOD

ans = pow(nCr(M, N) * factorial[N], 2, MOD)
for r in range(1, N + 1):
    p = nCr(M, r) * nPr(N, r) * pow(nPr(M - r, N - r), 2, MOD) % MOD
    p *= ((-1) ** ((r % 2) ^ 1))
    ans -= p
    ans %= MOD

print(ans)

感想

本番は  B の残りについての並べ方を考慮してなかった(なんで?????)
ちゃんと  {}_n P_r としてもっておきましょう。。。