【c++で楽しく実装!】二項係数(コンビネーション)の計算

zuka

二項係数の計算方法を解説するよ。

本記事では,二項係数の簡単な解説と実装例をお伝えしていきます。使用言語はc++です。その他の競技プログラミング関連の記事は,以下の目次をご覧ください。

目次

本記事の概要

まずは,二項係数の計算に関するお気持ちをお伝えしていきます。その次に,二項係数の素直な実装から始めて,徐々に発展的な実装をお伝えしていきます。

今回おさえるべき内容

 二項係数計算のお気持ち

 発展的な実装

お気持ち

二項係数の定義は以下の通りです。

\begin{align}
{}_n C_{k} &= \frac{n!}{k! (n-k)!} \label{basic}
\end{align}

まずは,この式を正直に実装すれば二項係数の計算はできてしまいます。しかし,競技プログラミングでは大きな素数で割ったあまりを考えるMODの世界に対応させてあげたり,より高速な実装方法が求められたりします。そこで,本記事では式($\ref{basic}$)を雛形として,いくつかの発展的な二項係数計算の実装方法をお伝えしていこうと思います。

なお,実装では以下のテンプレートを利用しています。

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define repi(i, a, b) for (int i = (int)(a); i < (int)(b); i++)
using namespace std;
typedef long long ll;

// 考える最大の整数
const int MAX = 500000;
// 今回採用する大きい素数
const int MOD = 1000000007;

ここでは,MODは大きな素数としてよく利用される1000000007を利用することにします。

雛形

式($\ref{basic}$)を正直に実装してみます。

ll nCk(int n, int k) {
    ll x = 1; // n!の初期値
    ll y = 1; // (n-k)!の初期値
    ll z = 1; // k!の初期値

    rep(i, n) x *= n - i; // n!を計算
    rep(i, n - k) y *= n - k - i; // (n-k)!を計算
    rep(i, k) z *= k - i; // k!を計算

    return (x / (y * z)); // 定義通りに計算
}

int main(){
  cout << nCk(5, 2) << endl; // 10が出力される
}
zuka

しっかり二項係数が計算できているね。long longを使うのはオーバーフロー対策だよ。

MODの世界への対応

大きな素数で割ったあまりを考えるMODの世界に対応するためには,単純に演算するごとにMODで割ったあまりを考えていけばOKです。c++では,計算結果が大きな値になりそうなかけ算では毎回大きな素数で割ってあげるクセを付けましょう。

ll nCk(int n, int k) {
    ll x = 1; // n!の初期値
    ll y = 1; // (n-k)!の初期値
    ll z = 1; // k!の初期値

    // 演算ごとにMODをとる
    rep(i, n) x = (x * (n - i)) % MOD; // n!を計算
    rep(i, n - k) y = (y * (n - k - i)) % MOD; // (n-k)!を計算
    rep(i, k) z = (z * (k - i)) % MOD; // k!を計算

    // 先に分母をMODの世界で計算してあげる
    ll yz = (y * z) % MOD;

    // 定義通りに計算
    return (x / (yz)) % MOD;
}

int main(){
  cout << nCk(5, 2) << endl; // 10が出力される
}

MODの世界の割り算

しかし!ここで問題が起きてしまいます。実は,MODの世界では割り算については少し気を付けなくてはならないことがあります。足し算・引き算・かけ算に関しては,MODを計算過程に組み入れることができました。

\begin{alignat}{3}
39 + 3 &\equiv 2 + 3 &\equiv 5 \quad &(\bmod 13) \\
39 – 3 &\equiv 3 – 3 &\equiv 0 \quad &(\bmod 13) \\
39 * 3 &\equiv 3 * 3 &\equiv 9 \quad &(\bmod 13) \\
\end{alignat}

39をあらかじめ法である13で割ってしまう操作ですね。法というのは,「$\bmod$ 〇〇」の〇〇の部分のことを指します。しかし,割り算だけは様子が異なります。

\begin{alignat}{3}
39 / 3 &\equiv 3 / 3 &\equiv 1 \quad &(\bmod 13) \\
\end{alignat}

zuka

この計算結果本当かな?

本来は,以下のような計算結果になるはずです。

\begin{alignat}{3}
39 / 3 &= 13 &\equiv 0 \quad &(\bmod 13) \\
\end{alignat}

証明は省略しますが,MODの世界の等式においては,両辺を割り算するときは法と互いに素な数である必要があります。この性質が,競技プログラミングでは少し厄介なのです。

今一度,先ほどの実装コードに立ち返ってみましょう。xyzをMODの世界で計算した後,割り算を行ってしまっています。これでは,数が十分大きいときには正しい答えが得られません。

ll nCk(int n, int k) {
    ll x = 1; // n!の初期値
    ll y = 1; // (n-k)!の初期値
    ll z = 1; // k!の初期値

    // 演算ごとにMODをとる
    rep(i, n) x = (x * (n - i)) % MOD; // n!を計算
    rep(i, n - k) y = (y * (n - k - i)) % MOD; // (n-k)!を計算
    rep(i, k) z = (z * (k - i)) % MOD; // k!を計算

    // 先に分母をMODの世界で計算してあげる
    ll yz = (y * z) % MOD;

    // 定義通りに計算
    return (x / (yz)) % MOD;
}

int main(){
  cout << nCk(5, 2) << endl; // 10が出力される
  cout << nCk(10000, 300) << endl; // 0が出力される <-違う!!(本当は634010422)
}

割り算をかけ算に直すという発想

MODの世界で割り算が行いにくいという問題を解決するために利用されるのが,フェルマーの小定理です。証明はWikipediaなどを参照してください。

フェルマーの小定理

$p$を素数,$a$ を$p$と互いに素な整数としたとき

\begin{align}
a^{p-1} \equiv 1 \quad(\bmod p) \label{Fermat}
\end{align}

式($\ref{Fermat}$)を利用すれば,ある整数の逆元を素数$p$のMODの世界で求めることができるようになります。逆元というのは,簡単には逆数のことです。しかし,MODの世界で考えると,逆元は逆数とは異なります。以下の例をみてください。

\begin{align}
1^{-1} &\equiv 1 \quad (\bmod 7)\\
2^{-1} &\equiv 4 \quad (\bmod 7)\\
3^{-1} &\equiv 5 \quad (\bmod 7)\\
4^{-1} &\equiv 2 \quad (\bmod 7)\\
5^{-1} &\equiv 3 \quad (\bmod 7)\\
\end{align}

この式の意味が分かりづらいですよね。これは,逆元の定義を踏まえると理解しやすいです。

逆元

\begin{alignat}{2}
ab &= ba &= 1
\end{alignat}

のとき,$a$は$b$の逆元であるという。

これをMODの世界で考えてみます。以下の等式を例にとります。

\begin{align}
2^{-1} &\equiv 4 \quad (\bmod 7)\\
\end{align}

実際に,逆元の定義に当てはめてみます。

\begin{align}
2^{-1} * 2 &= 4 * 2 \equiv 1 \quad (\bmod 7)\\
\end{align}

演算結果が1になっているため,7を法とするMODの世界では4は2の逆元であるといえます。


さて,逆数が求められればその値をかけ合わせることで割り算を行うことができます。式($\ref{Fermat}$)を変形すれば

\begin{align}
a \cdot a^{p-2} \equiv 1 \quad(\bmod p)
\end{align}

のようになります。計算結果が1であることに注目してください。これは,素数$p$のMODの世界では$a^{p-2}$が$a$の逆元であることを示しています。それでは,フェルマーの小定理式($\ref{Fermat}$)を利用して,先ほどの間違えている実装を直してみましょう。

// フェルマーの小定理に基づく逆元の計算
ll inv(ll x) {
    ll res = 1;
    rep(i, MOD - 2) res = (res * x) % MOD;
    return res;
}

ll nCk(int n, int k) {
    ll x = 1; // n!の初期値
    ll y = 1; // (n-k)!の初期値
    ll z = 1; // k!の初期値

    // 演算ごとにMODをとる
    rep(i, n) x = (x * (n - i)) % MOD; // n!を計算
    rep(i, n - k) y = (y * (n - k - i)) % MOD; // (n-k)!を計算
    rep(i, k) z = (z * (k - i)) % MOD; // k!を計算

    // 先に分母をMODの世界で計算してあげる
    ll yz = (y * z) % MOD;

    // 定義通りに計算
    return (x / (yz)) % MOD;
}

int main(){
  cout << nCk(5, 2) << endl; // 10が出力される
  cout << nCk(10000, 300) << endl; // 634010422が出力される <-OK!!
}
zuka

値が大きくなっても正しい値が出力されているね。でもめちゃくちゃ遅いよ。

高速化

先ほどのコードを実行すると,計算するのに10秒くらいかかってしまいます。競技プログラミングでは致命的な遅さです。そこで,先ほどのコードに高速化の工夫をしてみましょう。ここでは,思いつきやすい高速化として「計算結果を先にメモしておく」という方法を考えてみたいと思います。

まずは,階乗の計算結果をメモしておく部分の実装です。以下のコードでは,init関数を作ってmain関数の冒頭で呼び出すことで計算することを想定しています。

// 階乗を保管するメモ
ll fact[MAX];

void init() {
    // 1はじまりインデックスにそろえる初期値
    fact[0] = 1;
    fact[1] = 1;
    repi(i, 2, MAX){
        // 階乗計算
        fact[i] = fact[i - 1] * i % MOD;
    }
}

続いて,逆元を求めておく部分の実装です。ここでは,ちょっとしたテクニックを利用します。まず,以下は必ず成り立つ等式です。割り算は切り捨て除算とします。

\begin{align}
p &= p/a \cdot a + p \% a \\
\end{align}

続いて,両辺を$p$に関してmodをとります。

\begin{align}
p/a \cdot a + p \% a &\equiv 0 \\
p/a + (p \% a) \cdot a^{-1} &\equiv 0 \\
(p \% a) \cdot a^{-1} &\equiv – p/a \\
a^{-1} &\equiv – (p \% a)^{-1} \cdot p / a
\end{align}

このように,$a$の逆元$a^{-1}$を求めることができます。ただし,この値は負になるので,MODを足してあげることで適切な範囲に修正してあげます。

// 逆元を保管するメモ
ll inv[MAX];

void init() {
    inv[0] = 1;
    inv[1] = 1;
    repi(i, 2, MAX){
        // 逆元計算
        inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
    }
}

一応,結果の確認もしておきましょう。

// 今回は素数として7を利用
const int MOD = 7;
int main(){
  init();
  cout << inv[1] << endl; // 1 <- OK!!
  cout << inv[2] << endl; // 4 <- OK!!
  cout << inv[3] << endl; // 5 <- OK!!
  cout << inv[4] << endl; // 2 <- OK!!
  cout << inv[5] << endl; // 3 <- OK!!
}

最後に,逆元の階乗もあらかじめ計算しておきます。

// 逆元の階乗を保管するメモ
ll fact[MAX];

void init() {
    inv_fact[0] = 1;
    inv_fact[1] = 1;
    repi(i, 2, MAX){
        // 逆元の階乗を計算
        inv_fact[i] = inv_fact[i - 1] * inv[i] % MOD;
    }
}

最終形

さて,以上の高速化を踏まえた二項係数計算の最終形をまとめておきます。せっかくなので,二項係数計算の例外処理も加えておきます。

#include <bits/stdc++.h>
#define _GLIBCXX_DEBUG
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rrep(i, n) for (int i = (int)(n); i > 0; i--)
#define repi(i, a, b) for (int i = (int)(a); i < (int)(b); i++)
#define rrepi(i, a, b) for (int i = (int)(a); i > (int)(b); i--)
using namespace std;
typedef long long ll;

// 考える整数の最大値
const int MAX = 500000;
// 今回採用する大きい素数
const int MOD = 1000000007;

// メモを保管する場所
ll fact[MAX], inv_fact[MAX], inv[MAX];

// メモを計算する
void init() {
    // 初期値設定と1はじまりインデックスに直す
    fact[0] = 1;
    fact[1] = 1;
    inv[0] = 1;
    inv[1] = 1;
    inv_fact[0] = 1;
    inv_fact[1] = 1;
    // メモの計算
    repi(i, 2, MAX){
        // 階乗
        fact[i] = fact[i - 1] * i % MOD;
        // 逆元
        inv[i] = MOD - inv[MOD%i] * (MOD / i) % MOD;
        // 逆元の階乗
        inv_fact[i] = inv_fact[i - 1] * inv[i] % MOD;
    }
}

// 二項係数の実体
ll nCk(int n, int k) {
    ll x = fact[n]; // n!の計算
    ll y = inv_fact[n-k]; // (n-k)!の計算
    ll z = inv_fact[k]; // k!の計算
    if (n < k) return 0; // 例外処理
    if (n < 0 || k < 0) return 0; // 例外処理
    return x * ((y * z) % MOD) % MOD; //二項係数の計算
}

int main(){
  init();
  cout << nCk(5, 2) << endl; // 10が出力される
  cout << nCk(10000, 300) << endl; // 634010422が出力される
}
zuka

早い。

参考記事

二項係数 mod 素数を高速に計算する方法
「1000000007 で割ったあまり」の求め方を総特集! 〜 逆元から離散対数まで 〜
1~nまでの逆元をO(n)で求める方法

よかったらシェアしてね!

コメント

コメントする

目次
閉じる