Ccmmutty logo
Commutty IT
23 min read

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

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

AtCoder Beginner Contest 285

5 完。E 問題読み違えて出遅れ(´・ω・`) 今回、F までは難しめだった気がする。ノーペナ 6 完なら黄パフォだったっぽいし、F 問題が実装に時間かかりそうだったから飛ばしてみたけど、普通に F 問題やった方がよかったかも…。

A - Edge Checker 2

まあ、点が 1515 個しかないので、ii 番の点と結ばれていて、番号が ii より小さいものを配列 parent[i]parent[i] とかに入れれば、parent[b]=aparent[b]=a かどうか判断すればよくなる。
ちなみに、図をよく見ると、(2i2i2i+12i+1 を子とする完全二分木になっているので) [b2]=a\left[ \displaystyle\frac{b}{2} \right] =a かどうかを見ればいいことがわかる。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

int main() {
    int a, b;
    vector<int> parent({0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7});

    cin >> a >> b;

    // 条件演算子を使って出力 (もちろん if で書いてもいい)
    cout << (a == parent[b] ? "Yes" : "No");
    // cout << (a == b / 2 ? "Yes" : "No");  // こっちでも OK (parent がいらなくなる)

    return 0;
}

B - Longest Uncommon Prefix

22 重ループで i=1,2,,N1i=1, 2, \cdots, N-1 それぞれについて SkS_kSk+iS_{k+i} を比較して、ll を求めればいいだけ。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

int main() {
    int n;
    string s;

    cin >> n >> s;

    for (int i = 1; i < n; ++i) {
        int k;
        for (k = 0; k + i < n; ++k) {
            if (s[k] == s[k + i]) break;
        }
        cout << k << '\n';
    }

    return 0;
}

C - abc285_brutmhyhiizp

ID の長さが S|S| のものだけを考えた場合、A00 とした、S|S| 桁の 2626 進数として考えることができる。
また、これにより、長さが ii である ID は、26i26^i 個あることがわかる。
なので、答えは、(SSS|S| 桁の 2626 進数としたときの値) +i=1S126i+1+ \displaystyle\sum_{i=1}^{|S|-1} 26^i+1 となる。(11-indexed なので 11 足す。)
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

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

    cin >> s;

    // S を |S| 桁の 26 進数としたときの値を求める
    for (int i = 0, sz = s.size(); i < sz; ++i) {
        ans *= 26;
        ans += s[i] - 'A';
    }

    // |S| - 1 桁までの ID の個数を求める
    ll num = 1, sum = 0;
    for (int i = 1, sz = s.size(); i < sz; ++i) {
        num *= 26;
        sum += num;
    }

    ans += sum + 1;

    /* 式をうまいことやると、こっちでまとめて求められることがわかる。
    for (int i = 0, sz = s.size(); i < sz; ++i) {
        ans *= 26;
        ans += s[i] - 'A' + 1;
    } */

    cout << ans;

    return 0;
}

D - Change Usernames

TiT_i と一致する SjS_j がなければ、いつ、SiS_iTiT_i に変えても問題ない。
TiT_i と一致する SjS_j があるユーザ ii を考える。
ユーザ ii はユーザ jj が名前を変えるまで、名前を変えられない。なので、TjT_j と一致する SkS_k があるかを考える。なければ、ユーザ ii もユーザ jj も名前を変えられる。あれば、また、TkT_k と一致する SlS_l があるか考えなくてはならない。
ここで、「ユーザ ii はユーザ jj が名前を変えるまで、名前を変えられない」ことを iji \rightarrow j で表すとする。また、「TiT_i と一致する SjS_j がない」ことを iOKi \rightarrow \text{OK} で表すとする。
すると、すべてのユーザのユーザ名を変更できるということは、すべての ii で、
iOKi \rightarrow \cdots \rightarrow \text{OK}
となっているということと一致する。逆に、iOKi \rightarrow \cdots \rightarrow \text{OK} とならないのは、
ikki \rightarrow \cdots \rightarrow k \rightarrow \cdots \rightarrow k
のように、ループが発生する場合である。
上の話は、iji \rightarrow j をグラフ GGii から jj への辺を張ることに置きかえると、「すべてのユーザのユーザ名を変更できる \Leftrightarrow GG に閉路がない」になる。
なので、そのようなグラフ GG を作り、閉路があるか確かめればいい。閉路があるかどうかは、グラフの出次数が 00 の頂点を削除していくアルゴリズム (最終的に頂点がすべて削除されれば閉路なし、されなければ閉路あり) などで判断できる。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

// DAG (= 閉路がない有向グラフ) かどうか判断
// 出次数が 0 の頂点を削除していくアルゴリズム
bool isDAG(const vector<vector<int>> &g) {
    int sz = g.size();
    vector<vector<int>> reverseG(sz);
    vector<int> out(sz);
    vector<bool> deleted(sz);
    stack<int> leaves;

    for (int u = 0; u < sz; ++u) {
        for (int v : g[u]) {
            reverseG[v].push_back(u);
        }
    }

    for (int i = 0; i < sz; ++i) {
        out[i] = g[i].size();
        if (!out[i]) leaves.push(i);
    }

    int v;
    while (!leaves.empty()) {
        v = leaves.top();
        leaves.pop();

        deleted[v] = true;

        for (int a : reverseG[v]) {
            --out[a];
            if (!out[a]) leaves.push(a);
        }
    }

    for (int i = 0; i < sz; ++i) {
        if (!deleted[i]) return false;
    }

    return true;
}

int main() {
    int n;

    cin >> n;

    vector<string> s(n), t(n);

    map<string, int> mapS;  // mapS[S_i] = i
    for (int i = 0; i < n; ++i) {
        cin >> s[i] >> t[i];
        mapS[s[i]] = i;
    }

    vector<vector<int>> g(n);

    for (int i = 0; i < n; ++i) {
        auto itr = mapS.find(t[i]);
        // T_i = S_j となる j があるなら
        if (itr != mapS.end()) {
            g[i].push_back(itr->second);  // i -> j の辺を張る
        }
    }

    // 閉路がないなら "Yes" あるなら "No"
    cout << (isDAG(g) ? "Yes" : "No");

    return 0;
}

E - Work or Rest

曜日がループしていて、生産量は、曜日ではなく、「休日」から何日離れているかに依存するので、少なくとも 11 つの曜日が「休日」であることから、曜日 11 を「休日」としても一般性を失わない。 すると、dp[i]dp[i] を (曜日 ii を「休日」としたときの曜日 ii までの生産量) として、
dp[i]=max1k<i(dp[k]+sum[ik1])dp[i] = \max_{1 \leq k < i} (dp[k]+sum[i-k-1])
で求められる。ただし、dp[0]=0dp[0]=0 で、sum[i]sum[i]ii 日連続「平日」のときの生産量とする。
答えは、max1iN(dp[i]+sum[Ni])\displaystyle\max_{1 \leq i \leq N} (dp[i]+sum[N-i])
DP のオーダーは、sumsum を前計算しておけば、状態数が O(N)O(N)、遷移数が各 O(N)O(N)、遷移ごとの計算が O(1)O(1) なので、合わせて O(N2)O(N^2)
また、
sum[i]=k=1i2Ak+k=i2+1iAik+1=k=1i2Ak+k=1i2Ak\begin{align*} sum[i]&=\sum_{k=1}^{\left\lfloor \frac{i}{2} \right\rfloor}A_k+\sum_{k=\left\lfloor \frac{i}{2} \right\rfloor + 1}^{i}A_{i-k+1} \\ &=\sum_{k=1}^{\left\lfloor \frac{i}{2} \right\rfloor}A_k+\sum_{k=1}^{\left\lceil \frac{i}{2} \right\rceil}A_k \end{align*}
なので、sumsum の前計算は、AA の累積和を前計算すれば O(N)O(N)、愚直にやっても O(N2)O(N^2)
↓ こんな感じ。(00-indexed なので、少し式は違うけど)
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    int n;

    cin >> n;

    vector<ll> a(n);
    for (int i = 0; i < n; ++i) cin >> a[i];

    vector<ll> sum(n);
    for (int i = 1; i < n; ++i) {
        // 解説にはないけど、これでも sum を求められる。
        // 1-indexed なら sum[i] = sum[i - 1] + A_{ceil(i / 2)}
        sum[i] = sum[i - 1] + a[(i - 1) / 2];
    }

    vector<ll> dp(n);
    for (int i = 1; i < n; ++i) {
        for (int j = 0; j < i; ++j) {
            dp[i] = max(dp[i], dp[j] + sum[i - j - 1]);
        }
    }

    ll ans = 0;
    // 1-indexed なので、sum[n - i] -> sum[n - i - 1]
    for (int i = 0; i < n; ++i) ans = max(ans, dp[i] + sum[n - i - 1]);

    cout << ans;

    return 0;
}

F - Substring of Sorted String

時間外で解けたけど、解説いろいろ見て改良。
SSll 文字目から rr 文字目までからなる文字列を S[l:r]S[l:r] とおくと、S[l:r]S[l:r]TT の部分文字列であるかどうかは、
  • SSll 文字目から rr 文字目までが辞書順で昇順になっている。
  • S[l:r]S[l:r] に含まれる辞書順最小の文字を cc、最大の文字を dd としたとき、cc より大きく dd より小さいすべての文字が、SSS[l:r]S[l:r] に同じ個数含まれている。
の両方を満たすかどうかと同値である。
SSii 文字目を SiS_i とおくと、
  • 11 つ目の条件は、セグメント木 stst を用意し、st[i]=bool (SiSi+1)st[i]=\text{bool }(S_i \leq S_{i+1}) を保持すれば、[l,r)[l,r) の区間 AND\text{AND} で判断できる。更新も、取得も、クエリ毎 O(logN)O(\log N)
  • 22 つ目の条件は、文字の種類ごとに、文字数の区間和を取れれば判断できる。文字の種類数を σ(=26)\sigma (=26) とおくと、セグメント木や BIT で、更新も、取得も、クエリ毎 O(σlogN)O(\sigma \log N)
セグメント木や BIT は AtCoder Library にもあります。
オーダーは O(σ(N+Q)logN)O(\sigma(N+Q) \log N)
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

template <typename T>  // 0-indexed
class BIT {
    int n;
    vector<T> data;

   public:
    BIT(int _n) : n(_n), data(n + 1, 0) {}

    void add(int a, T w) {
        for (int i = ++a; i <= n; i += i & -i) data[i] += w;
    }

    T sum(int a) {  // [0, a]
        T ret = 0;
        for (int i = ++a; i > 0; i -= i & -i) ret += data[i];
        return ret;
    }

    T sum(int l, int r) {  // [l, r)
        T ret = 0;
        for (int i = r; i > 0; i -= i & -i) ret += data[i];
        for (int i = l; i > 0; i -= i & -i) ret -= data[i];
        return ret;
    }
};

template <typename S, S (*op)(S, S), S (*e)()>
class SegmentTree {
    int n;
    vector<S> d;

   public:
    SegmentTree(int n) : SegmentTree(vector<S>(n, e())){};
    SegmentTree(const vector<S>& v) : n(v.size()), d(n << 1) {
        for (int i = 0; i < n; ++i) {
            d[n + i] = v[i];
        }
        for (int i = n - 1; i > 0; --i) {
            d[i] = op(d[i << 1], d[i << 1 | 1]);
        }
    }

    void set(int pos, S x) {  // 1 点更新
        pos += n;
        d[pos] = x;
        for (int i = pos >> 1; i; i >>= 1) {
            d[i] = op(d[i << 1], d[i << 1 | 1]);
        }
    }

    S get(int pos) { return d[pos + n]; }

    S fold(int l, int r) {  // 区間取得
        S retl = e(), retr = e();
        for (l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if (l & 1) retl = op(retl, d[l++]);
            if (r & 1) retr = op(d[--r], retr);
        }
        return op(retl, retr);
    }
};

using S = bool;
S op(S a, S b) { return a & b; }
S e() { return true; }

int main() {
    int n, q;
    string s;

    cin >> n >> s >> q;

    vector<BIT<int>> charSum(26, BIT<int>(n));  // 文字の種類ごとの累積和
    SegmentTree<S, op, e> sorted(n - 1);        // ソートされているか
    // 初期化
    for (int i = 0; i < n - 1; ++i) sorted.set(i, s[i] <= s[i + 1]);
    for (int i = 0; i < n; ++i) charSum[s[i] - 'a'].add(i, 1);

    for (int i = 0; i < q; ++i) {
        int t;
        cin >> t;

        if (t == 1) {
            int x;
            char c;
            cin >> x >> c;
            --x;

            // 累積和の更新
            charSum[s[x] - 'a'].add(x, -1);
            charSum[c - 'a'].add(x, 1);
            s[x] = c;

            // ソートされているかの更新
            if (x > 0) sorted.set(x - 1, s[x - 1] <= s[x]);
            if (x < n - 1) sorted.set(x, s[x] <= s[x + 1]);
        } else if (t == 2) {
            int l, r;
            cin >> l >> r;
            --l;

            if (!sorted.fold(l, r - 1)) {  // ソートされてないなら
                cout << "No\n";
                continue;
            }

            int first, last;
            bool flag = true;  // 2 つ目の条件を満たすか
            for (first = 0; first < 26;) {
                if (charSum[first++].sum(l, r)) break;
            }
            for (last = 25; last >= 0;) {
                if (charSum[last--].sum(l, r)) break;
            }
            for (int i = first; i <= last; ++i) {
                if (charSum[i].sum(l, r) != charSum[i].sum(0, n)) {
                    flag = false;
                    break;
                }
            }

            cout << (flag ? "Yes" : "No") << '\n';
        }
    }

    return 0;
}

G - Tatami

AtCoder Library Practice Contest の Maxflow に似ていたので、解けそうだったけどダメでした…。
公式解説を見た感じ、"最小流量制限つき" 最大流問題というのを解ければ、Maxflow とほぼ同じみたい…。
まあ、知らなかったので、それでも解ける可能性のあったユーザ解説 (ygussany さん) で解いてみました。ユーザ解説では、GG のコピー 22 つを利用しています (オリジナルは利用していない) が、私は、GG (オリジナル) と GG のコピー (11 つ) を利用したものとして表記します。
例えば、入力例 11 では、以下のように流量 11 の辺を張ります。(1 には何も辺を張らない。)

これで、ss から tt に流し、流量が (2 の数) + (? の数) と一致するか判断すれば OK (ただし、数はオリジナルのみの数)。
↓ こんな感じ。
#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;

// 隣接するマスを返す
vector<pair<long long, long long>> move(long long x, long long y){
    vector<pair<long long, long long>> ret;

    ret.emplace_back(x-1,y);
    ret.emplace_back(x+1,y);
    ret.emplace_back(x,y-1);
    ret.emplace_back(x,y+1);

    return ret;
}

int main() {
    int h, w;

    cin >> h >> w;

    vector<string> c(h);
    for (int i = 0; i < h; ++i) cin >> c[i];

    int cntnot1 = 0;
    // [h * w, h * w * 2) はコピーの頂点
    // s は h * w * 2、t は h * w * 2 + 1
    mf_graph<int> g(h * w * 2 + 2);
    for (int i = 0; i < h; ++i) {
        for (int j = 0; j < w; ++j) {
            if (c[i][j] != '1') {  // 1 ではないなら
                ++cntnot1;  // カウント
                if ((i ^ j) & 1) {  // i + j が奇数のマスなら
                    g.add_edge(h * w * 2, i * w + j, 1);  // s からの辺を張る
                    g.add_edge((h + i) * w + j, h * w * 2 + 1, 1);  // コピーは t への辺を張る
                    for (auto [x, y] : move(i, j)) {
                        if (!(0 <= x && x < h && 0 <= y && y < w)) continue;
                        if (c[x][y] != '1') {  // 隣接マスが 1 ではないなら
                            g.add_edge(i * w + j, x * w + y, 1);  // 隣接マスへの辺を張る
                        }
                    }

                    if (c[i][j] == '?') {  // ? なら
                        g.add_edge(i * w + j, (h + i) * w + j, 1);  // コピーの同じマスへの辺を張る
                    }
                } else {  // i + j が偶数のマスなら
                    g.add_edge(i * w + j, h * w * 2 + 1, 1);  // t への辺を張る
                    g.add_edge(h * w * 2, (h + i) * w + j, 1);  // コピーは s からの辺を張る
                    for (auto [x, y] : move(i, j)) {
                        if (!(0 <= x && x < h && 0 <= y && y < w)) continue;
                        if (c[x][y] != '1') {  // 隣接マスが 1 ではないなら
                            g.add_edge((h + i) * w + j, (h + x) * w + y, 1);  // コピーに隣接マスへの辺を張る
                        }
                    }

                    if (c[i][j] == '?') {
                        g.add_edge((h + i) * w + j, i * w + j, 1);  // コピーの同じマスからの辺を張る
                    }
                }
            }
        }
    }

    // 完全マッチングなら "Yes" 違うなら "No"
    cout << (g.flow(h * w * 2, h * w * 2 + 1) == cntnot1 ? "Yes" : "No");

    return 0;
}

Ex - Avoid Square Number

公式解説見ました。累積和の部分までは思いつけたけど、累積和から元の数列を求められるの完全に忘れてた…。
形式的べき級数を考えると、(1+x+x2+)(1+x+x^2+\cdots) とか (1+x2+x4+)(1+x^2+x^4+\cdots) が現れそうだと考えられる。多項式の掛け算は、次数を nn とすると、畳み込みだと O(nlogn)O(n \log n) かかるけど、今回みたいなやつは累積和で O(n)O(n) で求められる。
これを思いつくと、O(NmaxEi)O(N \max E_i) 程度で通せそうな気がしてくる。
累積和で求める方法は、
f(x)=a0+a1x+a2x2+f(x)=a_0+a_1 x+a_2 x^2+ \cdots とおくと、
f(x)(1+x+x2++xi)=a0+(a0+a1)x++(k=0iak)xi+(k=1i+1ak)xi+1+f(x)(1+x+x^2+\cdots+x^i)=a_0 + (a_0+a_1) x + \cdots + \left(\sum_{k=0}^i a_k\right) x^i + \left(\sum_{k=1}^{i+1} a_k\right) x^{i+1} + \cdots
なので、次数が ii までの範囲であれば、(1+x+x2++xi)(1+x+x^2+\cdots+x^i) を掛けたい場合、f(x)f(x) の係数の累積和を取るだけで OK。
同様に、(1+x2+x4++x2i)(1+x^2+x^4+\cdots+x^{2i}) を掛けたい場合も、
f(x)(1+x+x2++xi)=a0+a1x1++(k=0ia2k)x2i+(k=0ia2k+1)x2i+1+(k=1i+1a2k)x2i+(k=1i+1a2k+1)x2i+1\begin{align*} f(x)(1+x+x^2+\cdots+x^i) &= a_0 + a_1 x^1 + \cdots + \left(\sum_{k=0}^i a_{2k}\right) x^{2i} + \left(\sum_{k=0}^i a_{2k+1}\right) x^{2i+1} \\ &+ \left(\sum_{k=1}^{i+1} a_{2k}\right) x^{2i} + \left(\sum_{k=1}^{i+1} a_{2k+1}\right) x^{2i+1} \end{align*}
なので、次数が 2i2i までの範囲であれば、f(x)f(x) の係数の累積和を、次数が奇数ごと、次数が偶数ごとに取るだけで OK。
多項式の掛け算が O(N)O(N) なので、包除原理で解けそう。
NN 要素の特定の ii 個が平方数である数列の数は、
j=1K[xEj](1+x2+x4+)i(1+x+x2+)Ni\prod_{j=1}^{K} [x^{E_j}] (1+x^2+x^4+\cdots) ^i (1+x+x^2+\cdots)^{N-i}
で表せる (ただし、[xi]f[x^i]fffxix^i の係数を表す) ので、包除原理を用いて、解は、
i=0N(1)i(Ni)j=1K[xEj](1+x2+x4+)i(1+x+x2+)Ni\sum_{i=0}^N (-1)^i \binom{N}{i} \prod_{j=1}^{K} [x^{E_j}] (1+x^2+x^4+\cdots)^i (1+x+x^2+\cdots)^{N-i}
となる。なので、f(i)=(1+x2+x4+)i(1+x+x2+)Nif(i)=(1+x^2+x^4+\cdots)^i (1+x+x^2+\cdots)^{N-i} とおくと、各 0iN0 \leq i \leq N について、f(i)f(i) を求められればいい。
まず、f(0)=(1+x+x2+)Nf(0)=(1+x+x^2+\cdots)^N は、(必要なのは、次数が maxEi\max E_i 以下の部分なので、) 累積和を使って、O(NmaxEi)O(N \max E_i) で求められる。また、
f(i)=f(i1)(1+x2+x4+)(1+x+x2+)f(i)=f(i-1) \frac{(1+x^2+x^4+\cdots)}{(1+x+x^2+\cdots)}
である。(1+x+x2+)(1+x+x^2+\cdots) を掛けるには、累積和を取ればいいので、(1+x+x2+)(1+x+x^2+\cdots) を割るには、隣接項の差分を取ればいい。
(1+x2+x4+)(1+x^2+x^4+\cdots) を掛けるのが O(maxEi)O(\max E_i)、差分も O(maxEi)O(\max E_i) で取れるので、各 f(i)f(i) (i>0)(i>0) は、f(i1)f(i-1) から O(maxEi)O(\max E_i) で求められる。
なので、各 0iN0 \leq i \leq N について、f(i)f(i) を求めるオーダーは O(NmaxEi)O(N \max E_i)
f(i)f(i) が求まれば、j=1K[xEj]f(i)\displaystyle\prod_{j=1}^{K} [x^{E_j}] f(i)O(K)O(K) で求められるので、解
i=0N(1)i(Ni)j=1K[xEj]f(i)\sum_{i=0}^N (-1)^i \binom{N}{i} \prod_{j=1}^{K} [x^{E_j}] f(i)
は、O(N(K+maxEi))O(N (K+\max E_i)) で求められる。
↓ こんな感じ。
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

// 階乗など
template <const int mod>
class Factorial {
    int n;
    vector<long long> fac;
    vector<long long> inv;

    long long pow(long long n, long long m) {
        long long ans = 1;
        while (m) {
            if (m & 1) ans = ans * n % mod;
            n = n * n % mod;
            m >>= 1;
        }
        return ans;
    }

   public:
    Factorial(int _n) : n(max(0, _n)), fac(n + 1), inv(n + 1) {
        fac[0] = 1;
        inv[0] = 1;
        for (int i = 1; i <= n; ++i) {
            fac[i] = fac[i - 1] * i % mod;
        }
        inv[n] = pow(fac[n], mod - 2);
        for (int i = n; i > 1; --i) {
            inv[i - 1] = inv[i] * i % mod;
        }
    }

    long long getFac(int a) { return fac[a]; }

    long long getInv(int a) { return inv[a]; }

    long long P(int a, int b) {
        if (a < b) return 0;
        return fac[a] * inv[a - b] % mod;
    }

    long long C(int a, int b) {
        if (a < b) return 0;
        return fac[a] * inv[b] % mod * inv[a - b] % mod;
    }
};

// f * (1 + x + x^2 + …)
void mul1(vector<ll> &f, const int mod) {
    for (int i = 1, sz = f.size(); i < sz; ++i) (f[i] += f[i - 1]) %= mod;
}

// f / (1 + x + x^2 + …)
void div1(vector<ll> &f, const int mod) {
    for (int i = f.size() - 1; i > 0; --i) (f[i] += mod - f[i - 1]) %= mod;
}

// f * (1 + x^2 + x^4 + …)
void mul2(vector<ll> &f, const int mod) {
    for (int i = 2, sz = f.size(); i < sz; ++i) (f[i] += f[i - 2]) %= mod;
}

int main() {
    int n, k;
    const int mod = 1000000007;
    ll ans = 0;

    cin >> n >> k;

    vector<int> e(k);
    for (int i = 0; i < k; ++i) cin >> e[i];

    Factorial<mod> F(n);

    // f = (1 + x + x^2 + …) で初期化
    vector<ll> f(10001, 1);
    for (int i = 1; i < n; ++i) mul1(f, mod);

    // f = (1 + x + x^2 + …)^{n - i} * (1 + x^2 + x^4 + …)^i
    for (int i = 0; i <= n; ++i) {
        ll squareI = 1;  // 特定の i 個が平方数である数列の数
        for (int j = 0; j < k; ++j) {
            (squareI *= f[e[j]]) %= mod;
        }
        // 包除原理
        (ans += (i & 1 ? mod - 1 : 1) * F.C(n, i) % mod * squareI) %= mod;
        div1(f, mod);
        mul2(f, mod);
    }

    cout << ans;

    return 0;
}

Discussion

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