Ccmmutty logo
Commutty IT
23 min read

水色コーダーが ABC 全問解説する【ABC279】

https://cdn.magicode.io/media/notebox/blob_fCXlYuJ

トヨタシステムズプログラミングコンテスト2022(AtCoder Beginner Contest 279)

6 完!今回は 6 完で青パフォ以上確定なので、やや難くらい?
ただ、4 ペナしたのはもったいない(´・ω・`)

A - wwwvvvvvv

何笑てんねん(´・ω・`) v 11 個につき、答えに 11 加算、w 11 個につき、答えに 22 加算すれば OK。文字列と for ループが使えれば簡単。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

int main() {
    string s;
    int ans = 0;

    cin >> s;

    for (int i = 0; i < s.length(); ++i) {
        if (s[i] == 'v') {  // i 文字目が 'v' なら
            ++ans;     // 1 加算
        } else {            // そうでないなら (i 文字目が 'w' なら)
            ans += 2;  // 2 加算
        }
    }

    cout << ans;

    return 0;
}
(v の個数) + (w の個数) × 22 で解いたけど、(文字列の長さ) + (w の個数) で解いたほうが賢かったかも?

B - LOOKUP

S,T100|S|, |T| \leq 100 なので、SSii 文字目から (i+T1)(i+|T|-1) 文字目が TT と一致するかを i=1i=1 から順番に試せばいい。i+T1Si+|T|-1 \leq |S| を満たさなくなったら終了。
(00-indexed で考えるなら、i=0i=0 から順番に、i+T1S1i+|T|-1 \leq |S|-1 (i+TS)(\Leftrightarrow i+|T| \leq |S|) を満たさなくなるまで。)
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

int main() {
    string s, t;

    cin >> s >> t;

    // 0-indexed 版
    for (int i = 0; i + t.length() <= s.length(); ++i) {
        bool flag = true;  // i 文字目から (i + |T| - 1) 文字目が T と一致しているか

        // 1 文字ずつ一致しているか調べる
        for (int j = 0; j < t.length(); ++j) {
            if (s[i + j] != t[j]) {  // 1 文字でも一致していなければ
                flag = false;  // 一致していないことが確定
                break;         // 終了
            }
        }
        if (flag) {  // 全部一致していれば
            cout << "Yes";  // "Yes" を出力して
            return 0;       // 終了
        }
    }

    cout << "No";  // どの場合も一致しなければ "No"

    return 0;
}

C - RANDOM

SSTT をそれぞれ列ごとに分解して考えてみる。例えば、入出力 11 ならこう。

SSTT は、同じ形の列で構成されていて、各種類の列の数がすべて同じ数になっている。この場合、同じ形の列を同じ位置に移動すればいいだけなので、SS の列を並び替えて TT を作れる。
入出力 22 の場合はこう。

今度は逆に、同じ形の列で構成されていて、各種類の列の数がすべて同じ数になっているわけではない。この場合は、明らかに一致しない列が現れるので、SS の列を並び替えて TT を作れない。
ということで、SSTT が、同じ形の列で構成されていて、各種類の列の数がすべて同じ数になっているかを判断できればいい。これは、SSTT をそれぞれ列でソートして、同じになるかどうかで判断できる。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

int main() {
    int h, w;

    cin >> h >> w;

    // 列でソートしたいので、s[列][行] になるようにする (t も同様)
    vector<vector<char>> s(w, vector<char>(h)), t(w, vector<char>(h));

    string in;  // 入力用

    // 入力 (s[列][行] にする)
    for (int i = 0; i < h; ++i) {
        cin >> in;
        for (int j = 0; j < w; ++j) {
            s[j][i] = in[j];
        }
    }
    // 入力 (t[列][行] にする)
    for (int i = 0; i < h; ++i) {
        cin >> in;
        for (int j = 0; j < w; ++j) {
            t[j][i] = in[j];
        }
    }

    // それぞれソート
    sort(s.begin(), s.end());
    sort(t.begin(), t.end());

    // 一致していれば "Yes"、してなければ "No"
    cout << (s == t ? "Yes" : "No");

    return 0;
}

D - Freefall

凸関数なので、三分探索できるけど、二分探索しか持ってなかったので、差分で考えました。
差分は、
f(x)f(x1)=B+AxAx1=B+Ax1xx(x1)=BA1x(x1)(x1+x)\begin{align*} f(x) - f(x-1) &= B + \dfrac{A}{\sqrt{x}} - \dfrac{A}{\sqrt{x-1}} \\ &= B + A \cdot \dfrac{\sqrt{x-1} - \sqrt{x}}{\sqrt{x(x-1)}} \\ &= B - A \cdot \dfrac{1}{\sqrt{x(x-1)}(\sqrt{x-1} + \sqrt{x})} \end{align*}
だから、単調増加。なので、差分が 00 未満になるかどうかで二分探索できる。(まあコンテスト中は、こんな計算しないで、単調増加っぽいなーってだけでにぶたんしたけど)
あとは、にぶたん (二分探索) するだけ。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    ll a, b;

    cin >> a >> b;

    // にぶたん
    ll ok = 1, ng = a / b + 1;  // 区間 (1, a / b + 1) を探索
    for (ll mid = (ok + ng) / 2; abs(ok - ng) > 1; mid = (ok + ng) / 2) {
        // f(mid) - f(mid - 1) < 0 なら
        if (b + (double)a / sqrt(mid) - (double)a / sqrt(mid - 1) < 0)
            ok = mid;  // mid は条件を満たす
        else
            ng = mid;  // mid は条件を満たさない
    }

    // 誤差で引っかからないように、小数点以下を適当に長めに出力
    printf("%.10lf", (double)a / sqrt(ok) + b * (ok - 1));

    return 0;
}

E - Cheating Amidakuji

操作を 11 つ抜いた場合と、操作を抜かなかった場合を比較してみる。
入力例 11 で、操作 22 を抜かなかった場合と抜いた場合を見てみると、
操作 22 で入れ替わる数字が逆になるだけだとわかる。
つまり、すべての操作を行った数列から、操作 ii で入れ替わった数字を入れ替えれば、操作 ii を抜いたときの数列になる。なので、すべての操作を行いつつ、各操作で入れ替わる数字を保持しておけば、すべての操作を行ったあとの数列を用いて、各操作を抜いた場合の数列をそれぞれ O(1)O(1) で求められる。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    int n, m;

    cin >> n >> m;

    // 0-indexed にして、B_j = 0 を満たす j を考える
    vector<int> b(n);
    for (int i = 0; i < n; ++i) b[i] = i;

    // 操作 i で入れ替わる数字を記録しながら、すべての操作を実行
    vector<pair<int, int>> swaped(m);  // 操作 i で入れ替わった数字のペア
    for (int i = 0, a; i < m; ++i) {
        cin >> a;
        // 0-indexed にしたので、(a, a + 1) => (a - 1, a)
        swaped[i] = {b[a - 1], b[a]};  // 入れ替わる数字を記録
        swap(b[a - 1], b[a]);          // 入れ替え
    }

    vector<int> inv(n);  // 数字 i がある位置を求める
    for (int i = 0; i < n; ++i) {
        inv[b[i]] = i;
    }

    // S_i を求める
    for (int i = 0; i < m; ++i) {
        // 入れ替えた数字のどちらかが 1 なら、
        // 1 ではないほうの数字の位置が、操作 i を抜いたときの 1 の位置
        // そうでないなら、
        // 操作 i を抜いても抜かなくても 1 の位置は変わらない
        // ※ 0-indexed にしたので、実際は、1 ではなく 0
        //    位置も最後に 1 加算する
        if (!swaped[i].first) {
            cout << inv[swaped[i].second] + 1 << '\n';
        } else if (!swaped[i].second) {
            cout << inv[swaped[i].first] + 1 << '\n';
        } else {
            cout << inv[0] + 1 << '\n';
        }
    }

    return 0;
}

F - BOX

ボールが 1010010^{100} 個とかあるけど、実際使うのは (N+Q)(N+Q) 個程度。Union-Find でいい感じにボールを管理すれば OK。ざっくりタイプごとに考えると、
タイプ 1: union する (箱が空の場合に注意)
タイプ 2: union する (箱が空の場合に注意)
タイプ 3: XX があるグループの代表から箱を割り出す
を行えばいい。
union するとき用に、箱に含まれるボールのうちの 11 つを、それぞれの箱ごとに保持しておくとよさそう。箱に含まれるボールのうちの 11 つを更新する必要があるのは、箱が空になったとき (箱に含まれるボールはなくなる) と、空の箱にボールが追加されたときのみかな。
グループの代表から箱を割り出すのは、グループの代表を添え字として、値を箱の番号とした配列を持っておけばいい。グループの代表が更新された可能性があるときにその配列も更新すればいい。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

// Union-Find 用のクラス
class UnionFind {
    vector<int> p;

   public:
    UnionFind(int n) : p(n, -1) {}  // n は要素数

    void unite(int x, int y) {  // union 関数
        x = find(x);
        y = find(y);
        if (x == y) return;

        if (p[x] < p[y]) {
            p[x] += p[y];
            p[y] = x;
        } else {
            p[y] += p[x];
            p[x] = y;
        }
        return;
    }

    int find(int x) {  // find 関数 (代表を返す)
        if (p[x] < 0) return x;
        while (p[p[x]] >= 0) {
            p[x] = p[p[x]];
            x = p[x];
            if (p[x] < 0) return x;
        }
        return p[x];
    }

    // 今回は使わない
    bool isSame(int x, int y) { return find(x) == find(y); }

    // 今回は使わない
    int size(int x) { return -p[find(x)]; }
};

int main() {
    int n, q;

    cin >> n >> q;

    // 箱に含まれるボールのうちの 1 つを保持する (空なら -1)
    vector<int> box(n + 1);
    // 代表から箱を割り出す用
    vector<int> inv(n + q + 1);

    // 最初は箱 i にボール i のみが入っているので、
    // 箱 i に含まれるボールのうちの 1 つはボール i のみで、
    // ボール i はそれぞれ箱 i の代表
    for (int i = 1; i <= n; ++i) box[i] = i, inv[i] = i;

    int ball = n;            // ボールの個数
    UnionFind u(n + q + 1);  // ボールは最大で (n + q) 個くらい

    for (int i = 0, t, x, y; i < q; ++i) {
        cin >> t >> x;
        if (t == 1) {  // タイプ 1
            cin >> y;  // 追加で y を受け取る
            if (box[y] == -1) continue;  // 箱 y が空なら何もしない

            if (box[x] == -1) {  // 箱 x が空なら
                // 箱 y の要素は箱 x の要素になる
                box[x] = box[y];
                // 代表が変わったので新たな代表が箱 x を指すようにする
                inv[u.find(box[x])] = x;
            } else {  // 両方中身があるなら
                // 箱 x と箱 y を union
                u.unite(box[x], box[y]);
                // 代表が変わった可能性があるので新たな代表が箱 x を指すようにする
                inv[u.find(box[x])] = x;
            }
            box[y] = -1;  // 箱 y は空にする
        } else if (t == 2) {  // タイプ 2
            ++ball;  // ボールの個数を加算

            if (box[x] == -1) {  // 箱 x が空なら
                // ボール ball は箱 x の要素になる
                box[x] = ball;
                // 代表が変わったので新たな代表が箱 x を指すようにする
                inv[u.find(box[x])] = x;
            } else {  // 空でないなら
                // 箱 x とボール ball を union
                u.unite(box[x], ball);
                // 代表が変わった可能性があるので新たな代表が箱 x を指すようにする
                inv[u.find(ball)] = x;
            }
        } else {  // タイプ 3
            // x があるグループの代表から箱を割り出して出力
            cout << inv[u.find(x)] << '\n';
        }
    }

    return 0;
}

G - At Most 2 Colors

時間外だけど、自力で通せた。(*´ω`*) それを置いとく。
長さ ii ごとに、条件を満たす塗り方を考える。そのとき、後ろの K1K−1 マス (ただし、iK1i \leq K-1 なら全マス) の塗り方で場合分けする。以下の 33 つは、長さが ii で、条件を満たす塗り方の数である。
dp1[i]dp1[i] : 後ろの K1K−1 マスが 11
dp2cis[i]dp2cis[i] : 後ろの K1K−1 マスが 22 色、かつ、マス ii とマス i1i-1 が同じ色
dp2trans[i]dp2trans[i] : 後ろの K1K−1 マスが 22 色、かつ、マス ii とマス i1i-1 が異なる色
このように分けると、以下の式が成り立つ。(ただし、dp2[i]=dp2cis[i]+dp2trans[i]dp2[i]=dp2cis[i]+dp2trans[i])
dp1[i]={C(1iK1)dp1[iK+1]×C+dp2[iK+1]×2(KiN)dp2cis[i]={0(i=1)dp2[i1](2iK1)dp2[i1]dp2trans[iK+2]×2(KiN)dp2trans[i]={0(i=1)dp1[i1]×(C1)+dp2[i1](2iN)\begin{align*} dp1[i]&= \begin{cases} C & (1 \leq i \leq K-1) \\ dp1[i-K+1] \times C + dp2[i-K+1] \times 2 & (K \leq i \leq N) \end{cases} \\ dp2cis[i]&= \begin{cases} 0 & (i=1) \\ dp2[i-1] & (2 \leq i \leq K-1) \\ dp2[i-1] - dp2trans[i-K+2] \times 2 & (K \leq i \leq N) \end{cases} \\ dp2trans[i]&= \begin{cases} 0 & (i=1) \\ dp1[i-1] \times (C-1) +dp2[i-1] & (2 \leq i \leq N) \end{cases} \end{align*}
  • dp1[i]dp1[i](1iK1)(1 \leq i \leq K-1) のとき、
    全マスを同じ色で塗る塗り方で、CC 色あるので CC 通り。
  • dp1[i]dp1[i](KiN)(K \leq i \leq N) のとき、
    後ろの K1K-1 マスを同じ色で塗る塗り方なので、長さ iK+1i-K+1 の塗り方の数を利用する。以下の 22 つを合わせて (dp1[iK+1]×C+dp2[iK+1]×2dp1[i-K+1] \times C + dp2[i-K+1] \times 2) 通り。
    • 長さ iK+1i-K+1 の後ろの K1K-1 マスが 11 色なら、
      次の K1K-1 マスを同じ色で塗るとき、選べる色は CC 色あるので、dp1[iK+1]×Cdp1[i-K+1] \times C 通り。
    • 長さ iK+1i-K+1 の後ろの K1K-1 マスが 22 色なら、
      次の K1K-1 マスを同じ色で塗るとき、選べる色はその 22 色のどちらかのみなので、dp2[iK+1]×2dp2[i-K+1] \times 2 通り。
  • dp2cis[i]dp2cis[i](2iK1)(2 \leq i \leq K-1) のとき、
    マス ii とマス i1i-1 が同じ色なので、長さ i1i-1 のときに 22 色で塗られている必要がある。選べる色は 11 色のみなので、dp2[i1]dp2[i-1] 通り。
  • dp2cis[i]dp2cis[i](KiN)(K \leq i \leq N) のとき、
    同様に、長さ i1i-1 のときに 22 色で塗られている必要があり、選べる色は 11 色のみなので、dp2[i1]dp2[i-1] 通り。
    しかし、dp2[i1]dp2[i-1] には、「後ろの K2K-2 マスが同じ色で塗られ、かつ、マス iK+2i-K+2 とマス iK+1i-K+1 の色が異なる」パターンが存在し、通り数は、dp2trans[iK+2]dp2trans[i-K+2] と一致する。そこからさらに同じ色で塗ってしまうと、後ろの K1K-1 マスすべてが同じ色で塗られてしまうので、そのパターンを引くと、(dp2[i1]dp2trans[iK+2]dp2[i-1] - dp2trans[i-K+2]) 通り。
  • dp2trans[i]dp2trans[i](2iN)(2 \leq i \leq N) のとき、
    以下の 22 つを合わせて (dp1[i1]×(C1)+dp2[i1]dp1[i-1] \times (C-1) +dp2[i-1]) 通り。
    • 長さ i1i-1 の後ろの K1K-1 マスが 11 色なら、
      次の 11 マスを異なる色で塗るとき、選べる色は C1C-1 色あるので、dp1[i1]×(C1)dp1[i-1] \times (C-1) 通り。
    • 長さ i1i-1 の後ろの K1K-1 マスが 11 色なら、
      次の 11 マスを異なる色で塗るとき、選べる色は 11 色のみなので、dp2[i1]dp2[i-1] 通り。
式のとおりに順番に計算すれば、O(N)O(N)。答えは、dp1[N]+dp2cis[N]+dp2trans[N]dp1[N] + dp2cis[N] + dp2trans[N]
ただし、K=2K=2 のときは正しく動かないので別で考える。答えは、CNC^N
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

// 法 mod での累乗
long long pow(long long n, long long m, long long mod) {
    long long ans = 1;
    while (m) {
        if (m & 1) ans = ans * n % mod;
        n = n * n % mod;
        m >>= 1;
    }
    return ans;
}

int main() {
    int n, k, c;
    const int mod = 998244353;

    cin >> n >> k >> c;

    // K = 2 なら C の N 乗が答え
    if (k == 2) {
        cout << pow(c, n, mod);
        return 0;
    }

    vector<ll> dp1(n + 1), dp2cis(n + 1), dp2trans(n + 1);

    // 以下、式のとおり
    dp1[1] = c % mod;
    // K - 1 以下
    for (int i = 2; i <= k - 1; ++i) {
        dp1[i] = dp1[i - 1];  // dp1[i] = c % mod;
        dp2cis[i] = (dp2cis[i - 1] + dp2trans[i - 1]) % mod;
        dp2trans[i] = (dp1[i - 1] * (c - 1) + dp2cis[i - 1] + dp2trans[i - 1]) % mod;
    }
    // K 以上
    for (int i = k; i <= n; ++i) {
        dp1[i] = (dp1[i - k + 1] * c + (dp2cis[i - k + 1] + dp2trans[i - k + 1]) * 2) % mod;
        dp2cis[i] = (dp2cis[i - 1] + dp2trans[i - 1] + (mod - dp2trans[i - k + 2])) % mod;
        dp2trans[i] = (dp1[i - 1] * (c - 1) + dp2cis[i - 1] + dp2trans[i - 1]) % mod;
    }

    cout << (dp1[n] + dp2cis[n] + dp2trans[n]) % mod;

    return 0;
}

Ex - Sum of Prod of Min

Ex 問題は解説読むだけでも難しい…(知らない定理やアルゴリズムも多いし)。解説読んで解いたので、公式解説の解説みたいなものを書きます。
まず、i=1NSi=M\displaystyle\sum_{i=1}^N S_i=M を満たすすべての SSk=1Nmin(k,Sk)\displaystyle\prod_{k=1}^N \min (k,S_k) の総和は、
k=1Ni=1min(k,i)xi\prod_{k=1}^N \sum_{i=1}^{\infty} \min (k,i) x^i
xMx^M の係数に一致する (maspy さんのページが参考になりそう)。「ffxnx^n の係数」を [xn]f[x^n]f と表記して、\sum の部分を展開すると、解は、
[xM]k=1N(x+2x2++kxk+kxk+1+kxk+2+)[x^M]\prod_{k=1}^N (x+2x^2+ \cdots + kx^k+ kx^{k+1}+kx^{k+2}+ \cdots)
となり、公式解説の式が出てくる。ここから、公式解説に合わせるため、kkii に置き換える。
fi(x)=x+2x2++ixi+ixi+1+ixi+2+f_i(x)=x+2x^2+ \cdots + ix^i+ ix^{i+1}+ix^{i+2}+ \cdots とおくと、解は、[xM]i=1Nfi(x)[x^M]\displaystyle\prod_{i=1}^N f_i(x) で、
(1x)fi(x)=fi(x)xfi(x)=x+x2++xi\begin{align*} (1-x)f_i(x)&=f_i(x)-xf_i(x) \\ &=x+x^2+ \cdots + x^i \end{align*}
fi(x)=x+x2++xif'_i(x)=x+x^2+ \cdots + x^i (=(1x)fi(x))(=(1-x)f_i(x)) とおくと、
(1x)fi(x)=fi(x)xfi(x)=xxi+1=x(1xi)\begin{align*} (1-x)f'_i(x)&=f'_i(x)-xf'_i(x) \\ &=x- x^{i+1} \\ &=x(1-x^i) \end{align*}
よって、(1x)2fi(x)=x(1xi)(1-x)^2f_i(x)=x(1-x^i) から、
fi(x)=x(1xi)(1x)2f_i(x)=\frac{x(1-x^i)}{(1-x)^2}
なので、解は、
[xM]i=1Nfi(x)=[xM]i=1Nx(1xi)(1x)2=[xM]i=1Nx(1xi)i=1N(1x)2=[xM]xNi=1N(1xi)(1x)2N=[xMN]i=1N(1xi)(1x)2N\begin{align*} [x^M]\prod_{i=1}^N f_i(x)&=[x^M]\prod_{i=1}^N \frac{x(1-x^i)}{(1-x)^2} \\ &=[x^M]\frac{\prod_{i=1}^N x(1-x^i)}{\prod_{i=1}^N (1-x)^2} \\ &=[x^M]\frac{x^N \prod_{i=1}^N(1-x^i)}{(1-x)^{2N}} \\ &=[x^{M-N}]\frac{\prod_{i=1}^N(1-x^i)}{(1-x)^{2N}} \end{align*}
ここで、[xMN]xki=1N(1xi)=0[x^{M-N}]x^k \prod_{i=1}^N(1-x^i)=0 (k>MN)(k>M-N) で、M2NM \leq 2N から N+1>MNN+1>M-N なので、
[xMN]i=1N(1xi)=[xMN](1xN+1)i=1N(1xi)=[xMN]i=1N+1(1xi)\begin{align*} [x^{M-N}]\prod_{i=1}^N(1-x^i)&=[x^{M-N}](1-x^{N+1})\prod_{i=1}^N(1-x^i) \\ &=[x^{M-N}]\prod_{i=1}^{N+1}(1-x^i) \end{align*}
同様に、N+k>MNN+k>M-N (k2)(k \geq 2) なので、
[xMN]i=1N+1(1xi)=[xMN](i=N+2(1xN+1))i=1N+1(1xi)=[xMN]i=1(1xi)\begin{align*} [x^{M-N}]\prod_{i=1}^{N+1}(1-x^i)&=[x^{M-N}]\left(\prod_{i=N+2}^{\infty}(1-x^{N+1})\right)\prod_{i=1}^{N+1}(1-x^i) \\ &=[x^{M-N}]\prod_{i=1}^{\infty}(1-x^i) \end{align*}
よって、解は、
[xMN]i=1N(1xi)(1x)2N=[xMN]i=1(1xi)(1x)2N\begin{align*} [x^{M-N}]\frac{\prod_{i=1}^N(1-x^i)}{(1-x)^{2N}}=[x^{M-N}]\frac{\prod_{i=1}^{\infty}(1-x^i)}{(1-x)^{2N}} \end{align*}
さらに、オイラーの五角数定理 i=1(1xi)=n=(1)nxn(3n1)2\displaystyle\prod_{i=1}^{\infty}(1-x^i)=\displaystyle\sum_{n=-\infty}^{\infty}(-1)^n x^{\frac{n(3n-1)}{2}} を利用すると、
[xMN]i=1(1xi)(1x)2N=[xMN]n=(1)nxn(3n1)2(1x)2N\begin{align*} [x^{M-N}]\frac{\prod_{i=1}^{\infty}(1-x^i)}{(1-x)^{2N}}=[x^{M-N}]\frac{\sum_{n=-\infty}^{\infty}(-1)^n x^{\frac{n(3n-1)}{2}}}{(1-x)^{2N}} \end{align*}
で、(1x)n=k=0(k+n1n1)xk(1-x)^{-n}=\displaystyle\sum_{k=0}^{\infty} \dbinom{k+n-1}{n-1} x^k (マクローリン展開とか) も利用して、an=(1)na_n=(-1)^nen=n(3n1)2e_n=\frac{n(3n-1)}{2} とおくと、
[xMN]n=(1)nxn(3n1)2(1x)2N=[xMN](n=anxen)(k=0(k+2N12N1)xk)\begin{align*} [x^{M-N}]\frac{\sum_{n=-\infty}^{\infty}(-1)^n x^{\frac{n(3n-1)}{2}}}{(1-x)^{2N}}=[x^{M-N}]\left(\sum_{n=-\infty}^{\infty}a_n x^{e_n}\right) \left(\sum_{k=0}^{\infty} \dbinom{k+2N-1}{2N-1} x^k \right) \end{align*}
xMNx^{M-N} の部分だけ抽出すると、
[xMN](n=(1)nxen)(k=0(k+n1n1)xk)=[xMN]ekMNakxek((MNek)+2N12N1)xMNek=[xMN]ekMNak(M+Nek12N1)xMN=ekMNak(M+Nek12N1)\begin{align*} [x^{M-N}]\left(\sum_{n=-\infty}^{\infty}(-1)^n x^{e_n}\right) \left(\sum_{k=0}^{\infty} \dbinom{k+n-1}{n-1} x^k \right)&=[x^{M-N}]\sum_{e_k \leq M-N}a_k x^{e_k} \cdot \dbinom{(M-N-e_k)+2N-1}{2N-1} x^{M-N-e_k} \\ &=[x^{M-N}]\sum_{e_k \leq M-N}a_k \dbinom{M+N-e_k-1}{2N-1} x^{M-N} \\ &=\sum_{e_k \leq M-N}a_k \dbinom{M+N-e_k-1}{2N-1} \end{align*}
これで、公式解説と同じ式になった。
二項係数を求めるのは 11 回あたり O(1)O(1) で可能だが、前計算した階乗の値が必要で、O(N)O(N) の階乗を求めるのに O(N)O(N) かかる。
そこで、Lucas の定理 (ni,ki(n_i, k_in,kn, kpp 進数表記での ii ビット目の値 ))
(nk)i(niki)(modp)\dbinom{n}{k} \equiv \displaystyle\prod_i \dbinom {n_i}{k_i} \pmod p
を用いると、mod1\bmod-1 の階乗までしか必要なくなるのでの計算量が、O(mod)O(\bmod) に落ちる。ただし、二項係数を求めるオーダーは O(1)O(1) から O(logmodN)O(\log_{\bmod} N) になる。
ekMNe_k \leq M-N を満たす kk は、O(N)O(\sqrt N) 個程度なので、O(NlogmodN+mod)O(\sqrt N \log_{\bmod} N + \bmod) でこの問題を解ける。
↓ こんな感じ (式をそのまま実装しただけだけど…)。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

// Lucas の定理を用いた二項係数の計算
// fact は前計算した階乗、inv は前計算した階乗の逆数
ll lucas(ll n, ll k, const ll mod, const vector<ll> &fact, const vector<ll> &inv) {
    if (n < k) return 0;
    ll ret = 1;

    while (n) {
        ll ni = n % mod;
        ll ki = k % mod;
        if (ni < ki) return 0;
        ret *= (fact[ni] * (inv[ki] * inv[ni - ki] % mod)) % mod;
        ret %= mod;
        n /= mod;
        k /= mod;
    }

    return ret;
}

// 法 mod での累乗 (逆数計算用)
long long pow(long long n, long long m, long long mod) {
    long long ans = 1;
    while (m) {
        if (m & 1) ans = ans * n % mod;
        n = n * n % mod;
        m >>= 1;
    }
    return ans;
}

int main() {
    ll n, m;
    ll ans = 0;
    const int mod = 200003;

    cin >> n >> m;

    // 階乗と階乗の逆数
    vector<ll> fact(mod), inv(mod);

    // 階乗の計算
    fact[0] = 1;
    for (int i = 1; i < mod; ++i) fact[i] = fact[i - 1] * i % mod;

    // 階乗の逆数の計算
    inv[mod - 1] = pow(fact[mod - 1], mod - 2, mod);
    for (int i = mod - 1; i > 0; --i) inv[i - 1] = inv[i] * i % mod;

    // 以下、式のとおり

    // e_k <= M - N (k >= 0)
    for (ll i = 0; i * (3 * i - 1) / 2 <= m - n; ++i) {
        ll ek = i * (3 * i - 1) / 2;
        ans += (i & 1 ? mod - 1 : 1) * lucas(m + n - ek - 1, 2 * n - 1, mod, fact, inv);
        ans %= mod;
    }

    // e_k <= M - N (k < 0)
    for (ll i = -1; i * (3 * i - 1) / 2 <= m - n; --i) {
        ll ek = i * (3 * i - 1) / 2;
        ans += (i & 1 ? mod - 1 : 1) * lucas(m + n - ek - 1, 2 * n - 1, mod, fact, inv);
        ans %= mod;
    }

    cout << ans;

    return 0;
}

Discussion

コメントにはログインが必要です。