はじめに
ABC172のE - NEQで包除原理を勉強していたところ、かなり理解するのに苦労してしまったので自分の言葉で整理しておきたいと思い、この記事を書きました。同じように苦しんでいる方の一助になれば幸いです。
本題
問題文
$1$ 以上 $M$ 以下の整数からなる長さ $N$ の数列 $A_1,A_2,\cdots ,A_N$ と $B_1,B_2,\cdots ,B_N$ の組であって、以下の条件をすべて満たすものの個数を求めてください。
- $1\le i\le N$ なる任意の $i$ について $A_i\neq B_i$
- $1\le i\lt j\le N$ なる任意の $(i,j)$ について $A_i\neq A_j$ かつ $B_i\neq B_j$
ただし、答えは非常に大きくなる可能性があるので、$(10^9+7)$ で割ったあまりを出力して下さい。
制約
- $1\le N\le M\le 5\times 10^5$
- 入力はすべて整数
包除原理は以下が成立するという原理です。（$N=3$ の場合を考えてみると理解しやすいと思います.）
集合 $X$ の要素の個数を $|X|$ とすると集合の列 $X_1,X_2,\cdots ,X_N$ に対して以下が成立する.
$$\begin{align}\left |\ \bigcup_{i=1}^{N}\ {X_i}\ \right | & = \sum_i{|X_i|} - \sum_{i<j}|X_i \cap X_j| + \sum_{i<j<k}|X_i\cap X_j\cap X_k| - \cdots \\ &\qquad+(-1)^{N+1}|X_1\cap X_2\cap \cdots\cap X_N| \end{align}$$
原理をなるべくそのままの形で適用することを易しい説明とする方針で書いていきます。
一つ目の条件より、以降は $1$ 以上 $M$ 以下の相異なる要素からなる数列を良い数列とします。
二つ目の条件について考えます。良い数列の組 $(A,B)$ からなる全体集合を $U$ としたとき、集合の列 $X_i\subset U\ (i=1,2,\cdots ,N)$ を $X_i := \Bigl\lbrace (A,B)\ \Big| \ A_i=B_i \Bigr\rbrace$ という風に定義してみましょう。このようにすることで求める値は $$\left | \overline{X_1\cup X_2\cup X_3\cup \cdots \cup X_N} \right|$$ であることがわかります。少し変形していくと
$$\begin{align} \left | \overline{X_1\cup X_2\cup X_3\cup \cdots \cup X_N} \right| &= \big|U\big| - \left | X_1\cup X_2\cup X_3\cup \cdots \cup X_N \right| \\ &= ({}_M P_N )^2 - \left | X_1\cup X_2\cup X_3\cup \cdots \cup X_N \right|\end{align}$$ とできるので後ろの部分を包除原理で求めていきましょう
$$\begin{align}\left | X_1\cup X_2\cup X_3\cup \cdots \cup X_N \right | = & +\sum_i{|X_i|} \\ &- \sum_{i<j}|X_i \cap X_j| \\ &+ \sum_{i<j<k}|X_i\cap X_j\cap X_k| \\\ &- \cdots \\ &+(-1)^{N+1}|X_1\cap X_2\cap \cdots\cap X_N| \end{align}$$
$k$行目$\ (j=1,2,\cdots ,N)$ に$k$個選んだときの積集合を足し合わせている項を置いて考えると、$k$行目では${}_N C_k $個の項が足されていることがわかります。ここで、$X_1,X_2,\cdots ,X_N$から$k$個をどのように取り出しても対称性から積集合の要素の個数は一定になります。つまり、$k$行目ではその行で固有のある定数を${}_N C_k $回足していることになります！これは$k$の値のみに依存するため$f(k)$とおきましょう。$\left(\left | X_1\cup X_2\cup X_3\cup \cdots \cup X_N \right | = \displaystyle\sum _{k=1}^{N}\ (-1)^{k+1}{}_N C_k \cdot f(k)\right)$ となっています。
$f(k)$ を求めましょう。良い数列 $A$ と $B$ で一致する場所がすでに $k$ 個選んであるので、そこに入れる整数の選び方が $ {}_M P_k $ 通り、残りの $N-k$ 個の場所に $A$ と $B$ それぞれにまだ入れていない整数を入れる方法の総数が $({}_{M-k}P_{N-k})^2$ 通りあるので $$f(k)={}_M P_k\ ({}_{M-k}P_{N-k})^2$$ とわかります。よって答えは $$({}_MP_N)^2 - \sum_{k=1}^N{(-1)^{k+1}{}_NC_k\cdot {}_MP_k\cdot ({}_{M-k}P_{N-k})^2}$$
とわかります。これは $0!,1!,2!,\cdots ,M!$ と $(0!)^{-1},(1!)^{-1},(2!)^{-1},\cdots ,(M!)^{-1}$ を $\pmod{10^9+7}$ で事前に計算しておくことで $O(M)$ で求まります。
ACコード（Python）
$\pmod{P}$ においてある整数ひとつの逆元を計算するのは $O(\log P)$ かかりますが、$1\le a\le N$ なる $a$ の逆元 $a^{-1}$ を列挙するのには $O(N)$ のアルゴリズムが存在します。
さいごに
- 僕が見つけた記事には包除原理を使うための集合を具体的に設定しているものがなかったので、それで困っている人に届けばと思います。
- 包除原理をもっと勉強していきます。