#1191464 (C++17) No.3769 Even Contraction

• 問題: No.3769 Even Contraction
• ユーザー: marc2825
• 提出日時: 2026-10-04 20:46:00
• 言語: C++17 (gcc 15.3.0 + boost 1.92.0 + ACL)
• 結果: WA
• 実行時間: -
• コード長: 8,972 bytes
• コンパイル時間: 2,069 ms
• コンパイル使用メモリ: 242,820 KB
• 実行使用メモリ: 1,298,972 KB
• 最終ジャッジ日時: 2026-10-16 20:57:09
• 合計ジャッジ時間: 17,080 ms
• ジャッジサーバーID （参考情報）: judge2_0 / judge1_0

ファイルパターン	結果
sample	AC * 3
other	AC * 8 WA * 4 MLE * 2 -- * 33

ソースコード

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

using pii = pair<int,int>;

struct Solver {
    int N;
    vector<int> val;
    vector<unordered_set<int>> G;
    vector<char> alive;

    Solver(
        int N_,
        const vector<pii>& edges,
        const vector<int>& A
    ) : N(N_), val(A), G(N_), alive(N_, 1) {

        for (auto [u,v] : edges) {
            G[u].insert(v);
            G[v].insert(u);
        }
    }

    // 0 : 00 or 22
    // 1 : 02
    // 2 : 11
    // -1: 操作不能
    int type(int u, int v) const {
        int x = val[u];
        int y = val[v];

        if (x == 1 && y == 1) {
            return 2;
        }

        if (
            (x == 0 && y == 2) ||
            (x == 2 && y == 0)
        ) {
            return 1;
        }

        if (
            (x == 0 && y == 0) ||
            (x == 2 && y == 2)
        ) {
            return 0;
        }

        return -1;
    }

    bool connected(int u, int v) const {
        if (!alive[u] || !alive[v]) {
            return false;
        }

        return G[u].count(v);
    }

    // u-v を縮約。
    // surviving vertex id を返す。
    int contract(int u, int v) {
        assert(connected(u,v));

        int nv = (val[u] + val[v]) / 2;

        G[u].erase(v);
        G[v].erase(u);

        // small-to-large
        if (G[u].size() < G[v].size()) {
            swap(u,v);
        }

        vector<int> tmp(
            G[v].begin(),
            G[v].end()
        );

        for (int x : tmp) {
            G[x].erase(v);
            G[x].insert(u);
            G[u].insert(x);
        }

        G[v].clear();
        alive[v] = 0;

        val[u] = nv;

        return u;
    }

    vector<pii> all_edges() const {
        vector<pii> ret;

        for (int v=0; v<N; v++) {
            if (!alive[v]) continue;

            for (int u : G[v]) {
                if (v < u) {
                    ret.push_back({v,u});
                }
            }
        }

        return ret;
    }

    // ========================================================
    // 固定フェーズ
    //
    // 例:
    // order = {0,1,2}
    //
    // 00/22 を可能な限り
    // -> 02 を可能な限り
    // -> 11 を可能な限り
    //
    // 一度終えたフェーズには戻らない
    // ========================================================

    int phase_strategy(
        const array<int,3>& order
    ) {
        int ans = 0;

        for (int want : order) {
            queue<pii> q;

            for (auto [u,v] : all_edges()) {
                if (type(u,v) == want) {
                    q.push({u,v});
                }
            }

            while (!q.empty()) {
                auto [u,v] = q.front();
                q.pop();

                if (!connected(u,v)) {
                    continue;
                }

                if (type(u,v) != want) {
                    continue;
                }

                int w = contract(u,v);
                ans++;

                // w の値が変わったので、
                // 新しく want 型になった辺を追加
                for (int x : G[w]) {
                    if (type(w,x) == want) {
                        q.push({w,x});
                    }
                }
            }
        }

        return ans;
    }

    // ========================================================
    // 常時優先順位付き貪欲
    //
    // 例:
    // order={1,0,2}
    //
    // 常に
    //   02
    //   > 00/22
    //   > 11
    //
    // の優先順位で縮約する。
    //
    // フェーズ制ではないので、
    // 後から02が生えれば再び02を優先する。
    // ========================================================

    int priority_strategy(
        const array<int,3>& order
    ) {
        array<queue<pii>,3> qs;

        auto push_edge =
            [&](int u, int v) {

                if (!connected(u,v)) {
                    return;
                }

                int t = type(u,v);

                if (t != -1) {
                    qs[t].push({u,v});
                }
            };

        for (auto [u,v] : all_edges()) {
            push_edge(u,v);
        }

        int ans = 0;

        while (true) {
            pii take = {-1,-1};

            for (int want : order) {
                while (!qs[want].empty()) {
                    auto [u,v] = qs[want].front();
                    qs[want].pop();

                    if (!connected(u,v)) {
                        continue;
                    }

                    int now = type(u,v);

                    if (now == -1) {
                        continue;
                    }

                    // キュー投入時から種類が変わった
                    if (now != want) {
                        qs[now].push({u,v});
                        continue;
                    }

                    take = {u,v};
                    break;
                }

                if (take.first != -1) {
                    break;
                }
            }

            if (take.first == -1) {
                break;
            }

            int w = contract(
                take.first,
                take.second
            );

            ans++;

            for (int x : G[w]) {
                push_edge(w,x);
            }
        }

        return ans;
    }

    // ========================================================
    // 葉貪欲
    //
    // 合法な「葉-親」辺のみを縮約。
    // 複数候補があれば order の優先順位。
    // ========================================================

    int leaf_strategy(
        const array<int,3>& order
    ) {
        array<queue<int>,3> qs;

        auto push_leaf =
            [&](int v) {

                if (!alive[v]) return;
                if (G[v].size() != 1) return;

                int u = *G[v].begin();

                int t = type(v,u);

                if (t != -1) {
                    qs[t].push(v);
                }
            };

        for (int v=0; v<N; v++) {
            push_leaf(v);
        }

        int ans = 0;

        while (true) {
            int leaf = -1;

            for (int want : order) {
                while (!qs[want].empty()) {
                    int v = qs[want].front();
                    qs[want].pop();

                    if (!alive[v]) continue;
                    if (G[v].size() != 1) continue;

                    int u = *G[v].begin();

                    int now = type(v,u);

                    if (now == -1) {
                        continue;
                    }

                    if (now != want) {
                        qs[now].push(v);
                        continue;
                    }

                    leaf = v;
                    break;
                }

                if (leaf != -1) {
                    break;
                }
            }

            if (leaf == -1) {
                break;
            }

            int p = *G[leaf].begin();

            int w = contract(leaf,p);
            ans++;

            // w 自身が葉になる可能性
            push_leaf(w);

            // w の値変更により、
            // 隣接する葉の合法性が変化する
            for (int x : G[w]) {
                push_leaf(x);
            }
        }

        return ans;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;

    vector<pii> edges;

    for (int i=0; i<N-1; i++) {
        int u,v;
        cin >> u >> v;

        --u;
        --v;

        edges.push_back({u,v});
    }

    vector<int> A(N);

    for (int &x : A) {
        cin >> x;
    }

    vector<array<int,3>> perms;

    array<int,3> p = {0,1,2};

    do {
        perms.push_back(p);
    } while (
        next_permutation(
            p.begin(),
            p.end()
        )
    );

    int ans = 0;

    // --------------------------------------------------------
    // 1. 3フェーズ固定：6通り
    // --------------------------------------------------------

    for (auto order : perms) {
        Solver s(N,edges,A);

        ans = max(
            ans,
            s.phase_strategy(order)
        );
    }

    // --------------------------------------------------------
    // 2. 常時優先順位付き貪欲：6通り
    // --------------------------------------------------------

    for (auto order : perms) {
        Solver s(N,edges,A);

        ans = max(
            ans,
            s.priority_strategy(order)
        );
    }

    // --------------------------------------------------------
    // 3. 葉貪欲：6通り
    // --------------------------------------------------------

    for (auto order : perms) {
        Solver s(N,edges,A);

        ans = max(
            ans,
            s.leaf_strategy(order)
        );
    }

    cout << ans << '\n';
}