【競プロ】組合せと剰余計算

🕒 2020/04/17 🔄 2023/05/01

ここでは、組合せの総数 ${}_n \mathrm{C}_k$ を $10^9+7$ で割った余りを求める方法について見ていきます。 $n,k$ が1つずつ与えられて、それに対して ${}_n \mathrm{C}_k$ を計算する状況を考えます。競プロ記事の一覧はこちら。

📘 目次

これまでの組合せと剰余計算

競プロの問題では、 $n$ 個から $k$ 個を選ぶ方法の総数 ${}_n\mathrm{C}_k$ を、 $10^9+7$ で割ってその余りを答える、ということがよくあります。この値は\[ {}_n\mathrm{C}_k=\frac{n!}{k!(n-k)!} \]なので、 $n,k$ が $20$ 以下くらいであれば、直接分母と分子を計算して割り算をし、 $10^9+7$ で割って求めることができます。しかし、もっと大きくなるとこの方法では求められません。

$n,k$ が10000程度であれば、【競プロ】パスカルの三角形と組合せで見た内容で求めることができます。次の関係式\[ {}_n\mathrm{C}_k = {}_{n-1}\mathrm{C}_{k-1}+{}_{n-1}\mathrm{C}_k \]を繰り返し用いることで、順番に求めていく方法です。

この方法では、 $n,k$ の制約はだいぶ緩みましたが、これよりもさらに大きな値に対して ${}_n\mathrm{C}_k$ を求めないといけないケースもあります。よく考えると、先ほど見たやり方は、 ${}_n\mathrm{C}_k$ を1つ出すために、他の組合せの総数も計算しないといけないため、無駄な計算をたくさんしています。できれば、 ${}_n\mathrm{C}_k$ を直接求められた方がいいですね。

\[ {}_n\mathrm{C}_k=\frac{n!}{k!(n-k)!} \]を、 $10^9+7$ で割った余りを求めるのにやっかいなのは、分母の存在です。ここの計算がなんとかできればいいのですが、ここで【競プロ】逆元と剰余演算で見た内容が使えるんですね。

組合せと剰余計算その1

さて、\[ {}_n\mathrm{C}_k=\frac{n!}{k!(n-k)!} \]を $10^9+7$ で割った余りの計算を考えていきます。 $n,k$ がともに $10^7$ 程度だとして考えていきます。また、 $p=10^9+7$ とおきます。この $p$ は素数です。

【競プロ】逆元と剰余演算で見たことから、 $\bmod p$ の世界では、 $k!(n-k)!$ で割ることは $\{k!(n-k)!\}^{p-2}$ を掛けることと同じになるのでした。なので、
\begin{eqnarray} {}_n\mathrm{C}_k &=& \frac{n!}{k!(n-k)!} \\[5pt] &\equiv& n! \{k!(n-k)!\}^{p-2} \pmod p \\[5pt] \end{eqnarray}となります。こうなれば、割り算がなくなって掛け算で表現できるようになったので、次のようなコードで求められるようになります。

#include <iostream>
using namespace std;

long long MOD = 1e9 + 7;

long long modPow(long long x, long long a) {
  if (a == 1) return x;
  if (a % 2) return (x * modPow(x, a - 1)) % MOD;
  long long t = modPow(x, a / 2);
  return (t * t) % MOD;
}

long long modInv(long long x) {
  return modPow(x, MOD - 2);
}

long long modFact(long long x) {
  long long ret = 1;
  for (long long i = 1; i <= x; i++) {
    ret = (ret * i) % MOD;
  }
  return ret;
}

long long modComb(long long n, long long k) {
  long long a, b, c, d;
  a = modFact(n);
  b = modFact(k);
  c = modFact(n - k);
  d = (b * c) % MOD;
  return (a * modInv(d)) % MOD;
}

int main() {
  long long n, k; cin >> n >> k;
  cout << modComb(n, k);
  return 0;
}

n, k を受け取って ${}_n\mathrm{C}_k$ を $10^9+7$ で割った余りを計算する modComb で答えを出力しています。 modComb では、 $n!\pmod p$ の結果を a に、 $k! \pmod p$ の結果を b に、 $(n-k)!\pmod p$ の結果を c に入れています。階乗は modFact で別途計算しています。

d に $k!(n-k)! \pmod p$ の結果を入れて、 a に d の逆元を掛けたものが modComb の計算結果です。逆元の計算は、【競プロ】逆元と剰余演算と同じ内容です。

こうして、 ${}_n\mathrm{C}_k$ を $10^9+7$ で割った余りが求められます。

組合せと剰余計算その2

先ほど、 ${}_n\mathrm{C}_k$ を $10^9+7$ で割った余りの計算を考えました。このときは $n,k$ がともに $10^7$ 程度だとしていましたが、 $n$ が $10^9$ 程度で $k$ が $10^7$ 程度の場合を考えてみましょう。

逆元が求められるようになった今となっては、同じように $n! \{k!(n-k)!\}^{p-2}$ を計算すればいいと思うかもしれませんが、 $n$ が $10^9$ 程度なら $n!$ の計算が間に合いません。 $10^9$ 回の掛け算では時間オーバーとなってしまいます。

そこで、少し式変形をする必要が出てきます。 $n!$ を $(n-k)!$ で割ると、結局残るのは $n-k+1$ から $n$ までの $k$ 個の積です。これを $k!$ で割ることを考えればいいですね。こうすると、分子も分母も $k$ 回程度の計算で求めることができます。

この場合でも計算できるように変更したものが次のコードです。

#include <iostream>
using namespace std;

long long MOD = 1e9 + 7;

long long modPow(long long x, long long a) {
  if (a == 1) return x;
  if (a % 2) return (x * modPow(x, a - 1)) % MOD;
  long long t = modPow(x, a / 2);
  return (t * t) % MOD;
}

long long modInv(long long x) {
  return modPow(x, MOD - 2);
}

long long modPerm(long long n, long long k) {
  long long ret = 1;
  for (long long i = 0; i < k; i++) {
    ret = (ret * (n - i)) % MOD;
  }
  return ret;
}

long long modComb(long long n, long long k) {
  long long a, b;
  a = modPerm(n, k);
  b = modPerm(k, k);
  return (a * modInv(b)) % MOD;
}

int main() {
  long long n, k; cin >> n >> k;
  cout << modComb(n, k);
  return 0;
}

modComb では a に $n-k+1$ から $n$ までの積を、b に $1$ から $k$ までの積を計算して $10^9+7$ で割った余りを入れています。

1つ目のやり方よりこちらの2つ目のやり方の方だと、 $n$ が $10^9$ 程度と大きい場合でも使うことができます。

ここで見た内容を利用すれば、AtCoder ABC 034 C - 経路や AtCoder ABC 156 D - Bouquet などに挑戦できるでしょう。