【AtCoder】ABC350 G - Mediator

問題

解法1 クエリ平方分割

もしクエリが1つしかない場合、以下の方法で求められる。

森の連結成分ごとに適当に根を取り、DFSで各頂点の親を求めると、各クエリについて、 と の親のいずれかが候補になるため、親を とすると、辺 と辺 が存在するかを、辺集合を std::set などで管理して判定すれば良い。

問題はDFSを毎回していると間に合わないことであるが、 回ごとに行うことにして、代わりに候補として と の親だけでなく、最後にDFSを行った後に追加した辺すべてについて見れば良い。

提出コード(1391ms)

#include "my_template.hpp"
#include "data_structure/unionfind.hpp"
using namespace std;

constexpr i64 MOD = 998244353;
constexpr int BS = 300;

void solve() {
    INT(N, Q);
    set<pair<int, int>> es1, es2;
    vector<int> X(Q + 1);
    int l = 0;
    vector<vector<int>> G(N);
    vector<int> seen(N), par(N);
    UnionFind uf(N);
    REP(q, Q) {
        I64(a, b, c);
        i64 A = 1 + ((a * (1 + X[l])) % MOD) % 2;
        i64 B = ((b * (1 + X[l])) % MOD) % N;
        i64 C = ((c * (1 + X[l])) % MOD) % N;
        if (A == 1) {
            es2.insert({B, C});
            es2.insert({C, B});
        } else {
            int ans = -1;
            if (uf.same(B, C)) {
                vector<int> candidates;
                if (par[B] != -1) candidates.push_back(par[B]);
                if (par[C] != -1) candidates.push_back(par[C]);
                FORE(cans, candidates) {
                    if (es1.count({cans, B}) > 0 and es1.count({cans, C}) > 0) {
                        ans = cans;
                        break;
                    }
                }
            } else {
                vector<pair<int, int>> ces2(ALL(es2));
                FORE(e, ces2) {
                    if (e.first == B or e.first == C) {
                        int cans = e.second;
                        if ((es2.count({cans, C}) > 0 or es1.count({cans, C}) > 0) and (es2.count({cans, B}) > 0 or es1.count({cans, B}) > 0)) {
                            ans = cans;
                            break;
                        }
                    } else if (e.second == B or e.second == C) {
                        int cans = e.first;
                        if ((es2.count({cans, C}) > 0 or es1.count({cans, C}) > 0) and (es2.count({cans, B}) > 0 or es1.count({cans, B}) > 0)) {
                            ans = cans;
                            break;
                        }
                    }
                }
            }
            l++;
            X[l] = ans + 1;
            print(X[l]);
        }
        if (q % BS == 0) {
            // update
            FORE(e, es2) {
                es1.insert(e);
                G[e.first].push_back(e.second);
                uf.merge(e.first, e.second);
            }
            es2.clear();
            seen.assign(N, 0);
            par.assign(N, -1);
            REP(i, N) {
                if (seen[i] == 0) {
                    auto rec = [&](auto f, int c, int p) -> void {
                        par[c] = p;
                        seen[c] = 1;
                        FORE(e, G[c]) {
                            if (p != e) {
                                f(f, e, c);
                            }
                        }
                    };
                    rec(rec, i, -1);
                }
            }
        }
    }
    return;
}

int main() {
    solve();
    return 0;
}

解法2 マージテク

判定部分は解法1で説明したクエリが1つの場合と同じ。毎回DFSをしていると間に合わないが、 と のうち連結成分数が小さい方が であるとすると、 を根と見てDFSをして、 の親を と更新する。

こうすることでマージテクによって計算量が抑えられる。

提出コード(76ms)

#include "my_template.hpp"
#include "data_structure/unionfind.hpp"
using namespace std;

constexpr i64 MOD = 998244353;

void solve() {
    INT(N, Q);
    i64 lans = 0;
    vector<set<int>> G(N);
    vector<int> par(N);
    UnionFind uf(N);
    REP(q, Q) {
        I64(a, b, c);
        i64 A = 1 + ((a * (1 + lans)) % MOD) % 2;
        i64 B = ((b * (1 + lans)) % MOD) % N;
        i64 C = ((c * (1 + lans)) % MOD) % N;
        if (A == 1) {
            auto rec = [&](auto f, int cur, int p) -> void {
                par[cur] = p;
                FORE(e, G[cur]) {
                    if (p != e) {
                        f(f, e, cur);
                    }
                }
            };
            if (uf.size(B) < uf.size(C)) {
                rec(rec, B, C);
            } else {
                rec(rec, C, B);
            }
            uf.merge(B, C);
            G[B].insert(C);
            G[C].insert(B);
        } else {
            // A == 2
            int ans = -1;
            vector<int> candidates;
            if (par[B] != -1) candidates.push_back(par[B]);
            if (par[C] != -1) candidates.push_back(par[C]);
            FORE(cans, candidates) {
                if (G[B].count(cans) > 0 and G[C].count(cans) > 0) {
                    ans = cans;
                    break;
                }
            }
            lans = ans + 1;
            print(lans);
        }
    }
    return;
}

int main() {
    solve();
    return 0;
}

感想

そもそも根付き木としてみたときの親しか候補にならないというのもはっきりとはわからず、LCAや集合の共通部分（bitset）を考えていた気がする（bitsetも平方分割するとできるらしいが）。マージテクは考えてもわからなかったが、多分上記の部分が詰められなかったことに原因がある気がする。

マージテクで実装すればかなり楽だったので、通したかったな…。

Link Cut Treeを使った方法もあとで学んでおきたい。