ABC021 D - 多重ループ
問題
https://atcoder.jp/contests/abc021/tasks/abc021_d
考察
問題を簡単にしてみます。$1 \leq a_1 < a_2 < ... < a_k < n$を満たす組み合わせを考えます。$1$以上$n$以下の整数を$k$個選び、全て異なる整数であれば上記の条件を満たすように並べることができます（「並べ替えれば条件を満たせる」系の考え方だと思って慣れる）。つまり、$1$以上$n$以下の整数を$k$個選ぶ組み合わせに等しいです 。これは${}_n C_k$です。
本来は$1 \leq a_1 \leq a_2 \leq ... \leq a_k \leq n$でした。これは、$1$以上$n$以下の整数を重複を許して$k$個選ぶ組み合わせに等しいです。
重複組み合わせを簡単に説明します。 重複組み合わせを理解するには仕切りを使った考え方が重要です。
桃太郎将軍は鬼退治のために軍隊を作りたいです。猿山、犬里、雉子ヶ谷にはそれぞれ、猿、犬、雉子が無限匹住んでいます。桃太郎将軍はこれらの土地から猿、犬、雉子を何匹かずつ雇って、軍隊を編成することにしました。しかし、きびだんごの在庫の問題で合計6匹までしか雇えません。桃太郎将軍は、軍隊を編成するメンバーの種類のパターンが何通りあるかを知りたいです。これを求めてください。ただし、雇わない種類が存在してもよいとします。
3種類の動物がいるので、2つの仕切りを使ってメンバーを3グループに分けてみます。
- ○○|○○|○○
- ○○|○|○○○
- ○|○○○○|○
- ○○○○○||○
例えば上記のような分け方があります。○で示したもののうち、左のグループから順に猿、犬、雉子になると決めます。
- 猿猿|犬犬|鳥鳥
- 猿猿|犬|鳥鳥鳥
- 猿|犬犬犬犬|鳥
- 猿猿猿猿猿||鳥
この「6つの○」と「2つの|」の並べ方が求める場合の数に一致します。 この場合の数は、まず8つの場所（「6つの○」と「2つの|」の合計）を用意しておき、仕切りの場所2つ（もしくは○の場所6つ）を選ぶ（選ばなかった場所には自動的にもう一方が入る）、として解け、答えは${}{8} C{2} = {}_{8} C_6 = 28$通りです。
一般化すると、$n$種類から重複を許して$k$個選ぶ組み合わせは${}_{n-1+k} C_k$です。
脱線したので元の問題に戻ると、$1$以上$n$以下の整数を$k$個選ぶ重複組み合わせは${}_{n-1+k} C_k$なので、これが答えです。
ところで、この二項係数を求めるのにもひと工夫必要です。
なので、
$k!^{-1}$ と $(n-k)!^{-1}$ が必要です。これを求めるにはフェルマーの小定理を使います。
$p$ を素数、$a$ を $p$ の倍数でない整数として $a^{p-1} \equiv 1 (mod, p)$ が成立する。（引用元）
よって、$a \times a^{p-2}\equiv 1 (mod, p)$ となるので、$a^{p-2}$ は $a$ の逆元です。 これは、繰り返し二乗法を用いて$O(\log p)$で求まります。
https://atcoder.jp/contests/abc021/submissions/8651510
なお、この問題では必要ないですが、複数の二項係数を高速で求めたい場面は多々あります。そんな時は逆元テーブル（$1!, 2!, ..., n!$の逆元）を予め求めておく必要があります。それぞれ求めると$O(n \log p)$かかりますが、最初に$(n!)^{-1}$を$O(\log p)$で求め、$(n-1)!^{-1} = (n!)^{-1} \times n$であることを利用して後ろから計算していくと$O(n + \log p)$で求まります。
https://atcoder.jp/contests/abc021/submissions/8651936