はじめに

paiza×Qiitaコラボキャンペーンの対象問題であるmod7占い(paiza Sランク)をC++で解いてみたので、解説していきます。

問題

mod7占いからpaizaの問題ページに飛ぶことができます。

あなたは今、「mod7占い」というサービスを始めようと考えています。

mod7占いとは、整数が書かれた複数のカードの中から3枚を選び、そこに書かれた整数の和が7で割り切れるかどうかで運勢を決めようというものです。カードは必ず異なる3枚を選ぶ必要があり、全てのカードには全て異なる数字が書かれています。

占いというからには、7で割り切れる組み合わせはそれなりに少なくする必要があります。そこで、適当な複数のカードに対して、カードに書かれた3つの整数を足した和が7で割り切れるような組合せがいくつあるかを計算するプログラムを作成してください。

制約については以下で与えます。
問題文を要約して形式的にすると以下のようになります。

問題文

長さ$N$の整数列$A=(A_1,...,A_N)$が与えられます。
$A_i+A_j+A_k$を7で割った余りが0となる相異なる添字の組$(i,j,k)$の総数を求めて下さい。

制約

・$1≦N≦10^5$
・$0≦A_i＜2^{32}$

解答

一番最初に思いつくのは単純に$(i,j,k)$の組を全探索する方法ですが$O(N^3)$となり間に合いません。なので別の方法を探して高速に処理したいです。

まず、$A_i+A_j+A_k$を７で割った余りは、$(A_iを7で割った余り)+(A_jを7で割った余り)+(A_kを7で割った余り)$と同じであることから、$A$の要素をすべて$A_iを7で割った余り$に置き換えてよいです。

更に、正整数を7で割った余りは0~6の7通りであり、正整数 $0≦i,j,k≦6$ について、$(i,j,k)$の組み合わせを全探索するときの計算量は$O(7^3)$となり、これは十分高速です。

なお、これは主客転倒という考え方です。$(i,j,k)$の組み合わせではなく、$A$の要素を7で割った余りの組み合わせを考えることで高速化が可能になりました。

よって、あらかじめ$mod_x:=配列Aの要素のうち,7で割った余りがxである要素の個数$という配列を用意しておくことで、$i+j+k$を7で割った余りが0であるとき、$i,j,k$の関係に応じて以下のように答えを加算していきます。

$mod_i×mod_j×mod_k$　($i≠j≠k$)

$_{mod_i}C_3$　($i=j=k$)

$mod_k×_{mod_i}C_2$　($i=j かつ j≠k$)

$mod_i×_{mod_j}C_2$　($j=k かつ k≠i$)

$mod_j×_{mod_k}C_2$　($k=i かつ i≠j$)

これは、例えば$i=j=k$のとき$mod_i$に寄与している$A$の要素の組み合わせの総数、$i≠j≠k$のとき$mod_i,mod_j,mod_k$にそれぞれ寄与する$A$の要素から1個ずつ選ぶ方法の総数であるので、それぞれコンビネーションと積で計算できることからわかります。

ここで、$i≠j≠k$でないとき、例えば$i=j=k$のとき$mod_i≦2$のとき$mod_i$に寄与しているAの要素を重複なく3つ選ぶことができないので、そのような組み合わせは存在しません。したがって、$_nC_r$を計算するcomb関数の実装では$n<r$のとき$_nC_r=0$であるという条件が必要であることに注意しましょう。

また、$(i,j,k)$の組み合わせに興味があるので、$i≦j≦k$として全探索することに注意して下さい。これは、$i,j,k$をすべて0から6で全探索してしまうと、例えば$(i,j,k)=(1,2,2)$と$(i,j,k)=(2,1,2)$という組み合わせに対して重複して答えを加算してしまうためです。

以上より、全体としては入力部分がボトルネックなので、$O(N)$でこの問題を解くことができました。

コンビネーションを計算する関数の実装

$_nC_r$が解答に必要であることが分かったので、$_nC_r$を計算する関数の実装についても触れておきます。

$_nC_r$には、以下のような重要な性質が成立します。

_nC_r=_{n-1}C_{r-1}+_{n-1}C_{r}

また、$_nC_n=_nC_0=1$であり、この問題では$_nC_r=0$ ($n<r$)として実装したいことを考慮すれば、あとはこれらの性質を利用して、再帰関数で高速に$_nC_r$の値を計算することができます。具体的な実装については以下の実装例を参照して下さい。

なお、$_nC_r$の定義通りに計算しても構いません。

実装例

以下、私がC++で解いた実装例を掲載します。

解答例(C++)

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

ll comb(ll n, ll r) {
    if (n==r || r==0) {
        return 1;
    }else if(n<r) {
        return 0;
    }else {
        return comb(n-1, r-1)+comb(n-1, r);
    }
}

int main() {
    ll n;
    cin >> n;
    vector<ll> a(n);
    vector<ll> mod(7);
    for(ll i=0; i<n; i++) {
        cin >> a[i];
        a[i]%=7;
        mod[a[i]]++;
    }

    ll ans=0;
    for(ll i=0; i<7; i++) {
        for(ll j=i; j<7; j++) {
            for(ll k=j; k<7; k++) {
                if((i+j+k)%7==0) {
                    if(i==j && j==k) {
                        ans+=comb(mod[i],3);
                    }else if(i==j) {
                        ans+=mod[k]*comb(mod[i],2);
                    }else if(j==k) {
                        ans+=mod[i]*comb(mod[j],2);
                    }else if(i==k) {
                        ans+=mod[j]*comb(mod[k],2); 
                    }else {
                        ans+=mod[i]*mod[j]*mod[k];
                    }
                }
            }
        }
    }
    
    cout << ans << endl;
}

あとがき

今回はpaiza Sランク問題のmod7占いという問題の解説を行いました。

計算量を意識することで、考える対象を入れ替えるというのは競技プログラミングにおいてもとても重要な考え方なので、身につけておきましょう。