NEQ [AtCoder Beginner Contest 172 E]

https://atcoder.jp/contests/abc172/tasks/abc172_e

前提知識

• 個数系包除原理
• 素数mod上での二項係数

解説

https://atcoder.jp/contests/abc172/submissions/14774888

もうこの手の問題をやりすぎて、何となく包除かな？という所からそれっぽいコードを書いて、サンプルが合ったので出したらACだった。
以下、試行手順を作って記載する。

難しい所はどこか

数列を二つ縦に並べたときに横も縦もダブってる数が無いようにするのが条件。
数列内でダブりが無いように数えるのはそれほど難しくはないが、数列間でのダブりを抑えるのが難しい。
適当に作ったときに2箇所ダブったり、3箇所ダブったり、NGの部分の個数が変わってしまう。
このようにNGの個数が変わるような問題はかなり難しそうに見えるが、包除原理に慣れると包除原理で解決できるアレだとなる。

包除原理

包除原理といえば、
n(AorBorC) = n(A) + n(B) + n(C) - n(A&B) - n(B&C) - n(C&A) + n(A&B&C)
というものであるが、
もし、「cnt[i] := i個の&で表現できる組み合わせ数」が高速に計算できれば、
n(AorBorC) = cnt[1] - cnt[2] + cnt[3]
のように個数に着目した包除原理に変換することができる。

これを今回の問題でも応用しよう。
cnt[same] := A[i]=B[i]である部分がsame個以上ある組合せの個数
こうすると、答えはcnt[0] - cnt[1] + cnt[2] - cnt[3] +...となる。

cntの計算

cnt[same]が計算できれば、もう答えが導けるが、これは高校の組合せ計算をする。
結論から言うとこれが組合せ。
com.aCb(N, same) * com.aPb(M, same) * com.aPb(M - same, N - same) * com.aPb(M - same, N - same)
分けて書く。

• com.aCb(N, same) : N要素のうちどのsame個の要素が同じか
• com.aPb(M, same) : [1,M]の数からsame個の同じ所に数を割り当てる
• com.aPb(M - same, N - same) : 同じ数に割り当てていない数を配列Aの同じじゃないところに割り当てる
• com.aPb(M - same, N - same) : 同じ数に割り当てていない数を配列Bの同じじゃないところに割り当てる

素数mod上での二項係数ライブラリを用意しておく必要がある。

int N, M;
Comb<mint, 1010101> com;
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> M;
    
    mint ans = 0;
    rep(same, 0, N + 1) {
        mint cnt = com.aCb(N, same) * com.aPb(M, same) * com.aPb(M - same, N - same) * com.aPb(M - same, N - same);
        if (same % 2 == 0) ans += cnt;
        else ans -= cnt;
    }
    cout << ans << endl;
}