abc156_dの解説 mintと競技プログラミングにおける組み合わせ数nCr、順列の場合数nPr

2020/05/21 14:05

AtCoder Beginner Contest 156のD問題 Bouquetを解く。必要となる素数で割ったあまりを取得するクラスと組み合わせ数の計算方法の解決

感謝の正拳突き１万回ならぬ、競プロサイトで 1 日 1AC(正解)するように頑張る。

今日はAtCoder Beginner Contest 156 の D 問題 Bouquetです。

プログラミングというよりは数学的な考察が必要な問題でした。また、素数で割った余りの計算方法を知っていないと答えが出せない問題で、競技プログラミング対策の勉強が必須の問題です。まず回答から。

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1e9 + 7;

class mint {
    long long x;
public:
    mint(long long x=0) : x((x % MOD + MOD) % MOD) {}
    mint operator-() const {
      return mint(-x);
    }
    mint& operator+=(const mint& a) {
        if ((x += a.x) >= MOD) x -= MOD;
        return *this;
    }
    mint& operator-=(const mint& a) {
        if ((x += MOD-a.x) >= MOD) x -= MOD;
        return *this;
    }
    mint& operator*=(const  mint& a) {
        (x *= a.x) %= MOD;
        return *this;
    }
    mint operator+(const mint& a) const {
        mint res(*this);
        return res+=a;
    }
    mint operator-(const mint& a) const {
        mint res(*this);
        return res-=a;
    }
    mint operator*(const mint& a) const {
        mint res(*this);
        return res*=a;
    }
    mint pow(long long t) const {
        if (!t) return 1;
        mint a = pow(t>>1);
        a *= a;
        if (t&1) a *= *this;
        return a;
    }
    // for prime MOD
    mint inv() const {
        return pow(MOD-2);
    }
    mint& operator/=(const mint& a) {
        return (*this) *= a.inv();
    }
    mint operator/(const mint& a) const {
        mint res(*this);
        return res/=a;
    }

    friend ostream& operator<<(ostream& os, const mint& m){
        os << m.x;
        return os;
    }
};

// 階乗を計算する
vector<mint> _factionals = {0,1};
mint factional(long long n) {
    if(_factionals.size() <= n) {
        long long  i = _factionals.size();
        _factionals.resize(n + 1);
        for(; i <= n; i++) {
            _factionals[i] = _factionals[i - 1] * i;
        }
    }
    return _factionals[n];
}

// 組み合わせ数を計算する。nCr
unordered_map<long long, vector<mint>> _combinations;
mint combination(long long n, long long r) {
    if(n < 1 || r < 1 || n < r) return mint(0);
    if(n == r) return mint(1);

    // nが小さい場合は公式を使う
    if(n < 51000) return factional(n) / (factional(r) * factional(n - r));

    // nが大きい場合はnC(r) = nC(r - 1) * (n - r + 1) / r
    // nC0 = 0, nC1 = nで初期化
    if(r > n - r) r = n - r;
    if(!_combinations.count(n)) {
        _combinations[n] = {0, n};
    }
    auto&& vec = _combinations[n];
    if(vec.size() <= r) {
        long long i = vec.size();
        vec.resize(r + 1);
        for(; i <= r; i++) vec[i] = (vec[i - 1] * (n - i + 1)) / i;
    }

    return vec[r];
}

// 順序のパターン数を計算する。nPr
unordered_map<long long, vector<mint>> _permutations;
mint permutation(long long n, long long r) {
    if(n < 1 || r < 1 || n < r) return mint(0);

    // nが小さい場合は公式を使う
    if(n < 51000) return factional(n) / factional(n - r);

    // nが大きい場合はnP(r) = nP(r - 1) * (n - r + 1)
    // nC0 = 0, nC1 = nで初期化
    if(!_permutations.count(n)) _permutations[n] = {0, n};
    auto&& vec = _permutations[n];
    if(vec.size() <= r) {
        long long i = vec.size();
        vec.resize(r + 1);
        for(; i <= r; i++) vec[i] = vec[i - 1] * (n - i + 1);
    }

    return vec[r];
}

int n,a,b;
mint result;
int main() {
    cin.tie(0); ios::sync_with_stdio(false); cout << fixed << std::setprecision(20);

    cin >> n >> a >> b;
    cout << mint(2).pow(n) - 1 - combination(n, a) - combination(n, b);
}

mint

MOD int の略らしい。ソースはこちらの(【競プロライブラリ】mint(modint)で学ぶ C++)[https://qiita.com/ue_sho/items/1ee5c3e665c72c035880]を参考にさせていただきました。

割った余りのとり方

競プロに親しんでいないと、割った余りの特性なんて考えたこともないと思います。こういう割った余りの特性を調べたりする数学を整数論と言い、(整数の合同)[https://ja.wikipedia.org/wiki/%E6%95%B4%E6%95%B0%E3%81%AE%E5%90%88%E5%90%8C]といったあたりのワードで語られます。

Wikipedia の記事がわかりやすく、2 時と 14 時が同じようなものを表すことは慣れ親しんだ数字ではないでしょうか。

足す、掛ける簡単です。足したり掛けたあとに MOD を取れば良いです。

int MOD = 12;
int tmp = 10;
tmp = (tmp + 4) % MOD; // 10時に4を足して、法が12時なので12で余りを取ると2時とわかる。
tmp = (tmp * 4) % MOD; // 時計で掛け算はしないが、同じ考え方

引き算は、マイナスになったら法を足せばよいです。

int MOD = 12;
int tmp = 2;
tmp = (tmp - 4); // 2時の4時間前が-2時になってしまう
if(tmp < MOD) tmp += MOD; // 法の12時を足すと、10時に戻る

割り算は解説記事がたくさんあるので省略しますが、普通に処理してしまうと誤った結果になります。

int MOD = 12;
int tmp = 17; // 本当は17時
tmp %= 12; // 5時にする
tmp / 5; // 1時 / 本当は割り切れないはずが1時か13時になってしまう

他に解説記事が多くあるので今回は省略しますが、フェルマーの小定理を使うとうまくいきます。

四則演算のルールが変わる

この〇〇で割った余りの世界では、四則演算もルールが変わったことがわかったかと思います。そのルールに則ったふるまいをしてくれるのが、mint クラスです。整数論を理解するのが王道ですが、理解に若干時間がかかる上に、理解していてもプログラムミスでバグにつながる場合もあります。必ず後で理解する必要はありますが、解けずに止まってしまうよりは mint クラスを使ったほうがよいのではないでしょうか。

べき乗

(AtCoder の公式解説 PDF)[https://img.atcoder.jp/abc156/editorial.pdf]にも載っている通り、べき乗は繰り返し２乗法で高速に計算可能です。mint クラスに実装されているので、利用してしまいましょう。

mint(2).pow(10); // 2の10乗で1024

組み合わせ

組み合わせの数の計算は高校数学でやりました。nCr の公式はn の階乗 / (r の階乗 × (n - r)の階乗です。ただし今回の問題は n <= 10 億となっていて、10 億の階乗を計算しているだけで TLE となってしまいます。

n が大きい場合(n > 1 * 10^6 くらい)

公式を使って解くことが先に思いつきますが、プログラミングの場合は漸化式を書いて計算したほうが計算量が少なくなる場合がありますす。組み合わせは次のようになります。 n / 1 // nC1 ↑ _ (n - 1) / 2 // nC2 ↑ _ (n - 2) / 3 // nC3 … ↑ * (n - r + 1) / r // nCr

今回の問題はプログラミングより数学的な整数の扱い方が中心でしたので、べき乗以外は計算量はさほど気にしなくても大丈夫でした。競技プログラミングで多くの場合、何回も組み合わせ数を計算する可能性が高いです。こういう場合を考えて、一度計算した値はメモしておくメモ化が有効です。

n が小さい場合

n が小さい場合は n の階乗を先に計算しておき、公式を使って計算したほうが高速です。

return factional(n) / (factional(r) * factional(n -r));

順列(abc156_d とは関係ない)

順列の場合の数も同様なので、この記事で紹介します。こちらも考え方は組み合わせと同様です。