#1101305 (C++23) No.3194 Do Optimize Your Solution

• 問題: No.3194 Do Optimize Your Solution
• ユーザー: 👑 tatyam
• 提出日時: 2025-06-27 02:58:12
• 言語: C++23 (gcc 13.3.0 + boost 1.87.0)
• 結果: MLE
• 実行時間: -
• コード長: 14,093 bytes
• コンパイル時間: 5,313 ms
• コンパイル使用メモリ: 344,784 KB
• 実行使用メモリ: 813,924 KB
• 最終ジャッジ日時: 2025-06-27 20:55:58
• 合計ジャッジ時間: 9,249 ms
• ジャッジサーバーID （参考情報）: judge4 / judge2

ファイルパターン	結果
sample	AC * 2
other	MLE * 1 -- * 16

ソースコード

#ifdef DEBUG
#include <range/v3/all.hpp>
#endif
#include <bits/stdc++.h>
using namespace std;
using u64 = uint64_t;
template<class T> using V = pmr::vector<T>;

/*
 木 A の辺ごとに寄与を数える．
 A を根付き木とする．
 A の頂点 v 以下の部分木に含まれる頂点の index の集合を S とする．B[S] に重み 1 を，それ以外の頂点に重み 0 を割り振ることを考える．
 
 以下のデータ構造があれば良い．
 
 add(v, w): 頂点 v に重み w を追加する．
 get(v): sum(dist(v, u) * w[u] for u in V) を求める．
 
 重心分解をすると，log 1 個になる
 
 HL 分解し，heavy-path を下から順に，
 */
u64 get(){
    u64 x = 0;
    uint8_t c;
    while ((c = getchar_unlocked() ^ 48) < 10){
        x = x * 10 + c;
    }
    return x;
}
void solve() {
    using namespace std::chrono;
    #ifdef DEBUG
    namespace views = ::ranges::views;
    namespace ranges = std::ranges;
    // stdin = fopen("/Users/tatyam/kyopro/c/c/test_in/01handmade007.txt", "r");
    #endif
    auto t0 = steady_clock::now();
    // 入力
    u64 N = get();
    auto t_inputN = steady_clock::now();
    cerr << "N: " << duration_cast<milliseconds>(t_inputN - t0).count() << "ms\n";
    // A の入力
    V<pair<u64, u64>> edge(N - 1);
    V<u64> deg(N);
    for (auto& [u, v] : edge) {
        u = get() - 1;
        v = get() - 1;
        deg[u]++;
        deg[v]++;
    }
    auto t_inputA = steady_clock::now();
    cerr << "A edge: " << duration_cast<milliseconds>(t_inputA - t_inputN).count() << "ms\n";
    V<V<u64>> A(N);
    for (auto&& [a, d] : views::zip(A, deg)) a.reserve(d);
    for (auto [u, v] : edge) {
        A[u].push_back(v);
        A[v].push_back(u);
    }
    auto t_inputA2 = steady_clock::now();
    cerr << "A build: " << duration_cast<milliseconds>(t_inputA2 - t_inputA).count() << "ms\n";
    auto t1 = steady_clock::now();
    cerr << "input: " << duration_cast<milliseconds>(t1 - t0).count() << "ms\n";
    // A を HL 分解
    // heavy-edge は G[*][0] に
    V<u64> siz(N, 1);
    siz.assign(N, 1);
    auto t_hld0 = steady_clock::now();
    auto hld = [&](auto hld, u64 v) -> void {
        if (A[v].empty()) return;
        for (u64 i : A[v]) {
            erase(A[i], v);
            hld(hld, i);
            siz[v] += siz[i];
        }
        auto p = ranges::max_element(A[v], {}, [&](u64 i) { return siz[i]; });
        swap(*p, A[v][0]);
    };
    hld(hld, 0);
    auto t_hld1 = steady_clock::now();
    cerr << "HL: " << duration_cast<milliseconds>(t_hld1 - t_hld0).count() << "ms\n";
    // A のオイラーツアー順に頂点番号の relabel
    V<u64> label(N, 0);
    auto t_relabel0 = steady_clock::now();
    {
        u64 idx = 0;
        auto relabel = [&](auto relabel, u64 v) -> void {
            label[v] = idx++;
            for (u64 i : A[v]) {
                relabel(relabel, i);
            }
        };
        relabel(relabel, 0);
    }
    auto t_relabel1 = steady_clock::now();
    cerr << "A relabel: " << duration_cast<milliseconds>(t_relabel1 - t_relabel0).count() << "ms\n";
    // A の作り直し
    auto t_Arebuild0 = steady_clock::now();
    {
        V<V<u64>> A_(N);
        for (auto&& [i, a] : A | views::enumerate) {
            for (u64& x : a) x = label[x];
            A_[label[i]] = std::move(a);
        }
        A = std::move(A_);
    }
    auto t_Arebuild1 = steady_clock::now();
    cerr << "A rebuild: " << duration_cast<milliseconds>(t_Arebuild1 - t_Arebuild0).count() << "ms\n";
    // heavy_parent: A における親の番号 (light-edge なら -1)
    auto t_parent0 = steady_clock::now();
    V<u64> heavy_parent(N, -1);
    for (auto&& [i, a] : views::enumerate(A)) {
        if (a.empty() || i == 0) continue; // 根はやらなくていいので切る
        heavy_parent[a[0]] = i;
    }
    auto t_parent1 = steady_clock::now();
    cerr << "A parent: " << duration_cast<milliseconds>(t_parent1 - t_parent0).count() << "ms\n";
    auto t3 = steady_clock::now();
    cerr << "A relabel+parent: " << duration_cast<milliseconds>(t3 - t1).count() << "ms\n";
    // B を入力
    auto t_Binput0 = steady_clock::now();
    deg.assign(N, 0);
    for (auto& [u, v] : edge) {
        u = label[get() - 1];
        v = label[get() - 1];
        deg[u]++;
        deg[v]++;
    }
    auto t_Binput1 = steady_clock::now();
    cerr << "B edge: " << duration_cast<milliseconds>(t_Binput1 - t_Binput0).count() << "ms\n";
    V<V<u64>> B(N);
    for (auto&& [a, d] : views::zip(B, deg)) a.reserve(d);
    for (auto [u, v] : edge) {
        B[u].push_back(v);
        B[v].push_back(u);
    }
    auto t_Binput2 = steady_clock::now();
    cerr << "B build: " << duration_cast<milliseconds>(t_Binput2 - t_Binput1).count() << "ms\n";
    auto t4 = steady_clock::now();
    cerr << "B input: " << duration_cast<milliseconds>(t4 - t3).count() << "ms\n";
    // B を重心分解
    struct T {
        u64 cnt = 0, sum = 0;
        void operator+=(T a) {
            cnt += a.cnt;
            sum += a.sum;
        }
        void operator-=(T a) {
            cnt -= a.cnt;
            sum -= a.sum;
        }
    };
    V sum_dist(N, array<T, 2>{}); // { (重心が i である連結成分) -> i, (重心が i である連結成分) -> parent[i] }
    struct C { array<T, 2>* cen; u64 dist; };
    V<V<C>> centroid_ancestor(N); // {sum_dist の index，重心までの距離}
    for (auto& c : centroid_ancestor) c.reserve(bit_width(N));
    auto t_centroid0 = steady_clock::now();
    {
        V<bool> deleted(N, false);
        V<pair<u64, u64>> cc; // {index, 部分木サイズ}
        auto adj = [&](u64 v) {
            return B[v] | views::filter([&](u64 u) { return !deleted[u]; });
        };
        
        auto CD = [&](auto CD, u64 v, u64 siz_cc) -> u64 {
            // 頂点 v を含む連結成分を重心分解する
            const u64 half_siz = siz_cc / 2;
            v = ranges::min(cc, {}, [&](const auto& p) { return p.second >= half_siz ? p.second : u64(-1); }).first;
            // 重心を削除
            deleted[v] = true;
            erase_if(B[v], [&](u64 u) { return deleted[u]; });
            // centroid_ancestor を更新しながら，子の連結成分の siz を求める
            const u64 cen = v;
            centroid_ancestor[v].emplace_back(bit_cast<array<T, 2>*>(cen), 0);
            // 子を重心分解して，B を重心分解木に変更
            for (u64& u : B[v]) {
                cc.clear();
                auto dfs = [&](auto dfs, u64 parent, u64 v, u64 dist) -> u64 {
                    centroid_ancestor[v].emplace_back(bit_cast<array<T, 2>*>(cen), dist);
                    u64 siz = 1;
                    for (u64 u : adj(v)) if  (u != parent) {
                        siz += dfs(dfs, v, u, dist + 1);
                    }
                    cc.emplace_back(v, siz);
                    return siz;
                };
                dfs(dfs, v, u, 1);
                u = CD(CD, u, siz[u]);
            }
            return v;
        };
        {
            auto dfs = [&](auto dfs, u64 parent, u64 v) -> u64 {
                u64 siz = 1;
                for (u64 u : adj(v)) if  (u != parent) {
                    siz += dfs(dfs, v, u);
                }
                cc.emplace_back(v, siz);
                return siz;
            };
            dfs(dfs, -1, 0);
        }
        const u64 cen = CD(CD, 0, N);
        auto t_centroid1 = steady_clock::now();
        cerr << "centroid decomp: " << duration_cast<milliseconds>(t_centroid1 - t_centroid0).count() << "ms\n";
        // 重心分解木 B の relabel
        label.assign(N, -1);
        u64 idx = 0;
        V<u64> sizes;
        sizes.reserve(N);
        auto t_Brelabel0 = steady_clock::now();
        auto relabel = [&](auto relabel, u64 v) -> void {
            // v の部分木で，まだ label のついていない頂点を relabel
            // siz が上位 √n 個の頂点を取り出して再帰，残りの部分木 (√n 頂点以下) もそれぞれ再帰 => 再帰的な平方分割
            // siz を計算しながら，上位 √n 個のボーダーを求める
            sizes.clear();
            auto& q = deg; // BFS 順 (メモリ使い回し)
            q.clear();
            q.push_back(v);
            for (auto p = q.begin(); p != q.end(); ++p) {
                u64 v = *p;
                for (u64 u : B[v]) q.push_back(u);
            }
            for (u64 v : q | views::reverse) {
                auto& s = siz[v];
                s = u64(1) << 32;
                for (u64 u : B[v]) s += siz[u];
                s &= u64(-1) << 32;
                s |= v; // 頂点番号を付けて tie-break
                sizes.push_back(s);
            }
            if (sizes.size() <= 1) {
                // 1 個しかないなら，そのまま relabel
                label[v] = idx++;
                return;
            }
            const u64 sq = sqrt(sizes.size()); // 上位 sq 個を取り出す
            const u64 lb = [&] {
                auto p = sizes.begin() + (sq - 1);
                ranges::nth_element(sizes, p, greater<>{});
                return *p;
            }();
            // siz が lb 未満になったら再帰，辺を削除
            auto dfs = [&](auto dfs, u64 v) -> void {
                erase_if(B[v], [&](u64 u) {
                    if (siz[u] < lb) {
                        relabel(relabel, u);
                        return 1;
                    } else {
                        dfs(dfs, u);
                        return 0;
                    }
                });
            };
            dfs(dfs, v);
            return relabel(relabel, v);
        };
        relabel(relabel, cen);
        auto t_Brelabel1 = steady_clock::now();
        cerr << "B relabel: " << duration_cast<milliseconds>(t_Brelabel1 - t_Brelabel0).count() << "ms\n";
        // centroid_ancestor を更新
        auto t_centroidAnc0 = steady_clock::now();
        for (auto& anc : centroid_ancestor) {
            for (auto& [cen, dist] : anc) {
                cen = &sum_dist[label[bit_cast<u64>(cen)]];
            }
        }
        auto t_centroidAnc1 = steady_clock::now();
        cerr << "centroid_ancestor update: " << duration_cast<milliseconds>(t_centroidAnc1 - t_centroidAnc0).count() << "ms\n";
    }
    auto t5 = steady_clock::now();
    cerr << "centroid: " << duration_cast<milliseconds>(t5 - t4).count() << "ms\n";
    // B の全頂点に 0 をセット
    auto t_Bzero0 = steady_clock::now();
    for (auto& anc : centroid_ancestor) {
        T D;
        {
            auto [cen, dist] = anc[0];
            auto& [P, Q] = *cen;
            D = {1, dist};
            P += D;
        }
        for (auto [cen, dist] : anc | views::drop(1)) {
            auto& [P, Q] = *cen;
            Q += D;
            D = {1, dist};
            P += D;
        }
    }
    auto t_Bzero1 = steady_clock::now();
    cerr << "B zero: " << duration_cast<milliseconds>(t_Bzero1 - t_Bzero0).count() << "ms\n";
    u64 ans = 0, sum01 = 0;
    auto set = [&](u64 v) -> void {
        // B で頂点 v の重みを 1 に
        auto& anc = centroid_ancestor[v];
        T D;
        {
            auto [cen, dist] = anc[0];
            auto& [P, Q] = *cen;
            D = {1, dist};
            P -= D;
            sum01 += P.sum + P.cnt * D.sum;
            P -= D;
        }
        for (auto [cen, dist] : anc | views::drop(1)) {
            auto& [P, Q] = *cen;
            Q -= D;
            sum01 -= Q.sum + Q.cnt * D.sum;
            Q -= D;
            D = {1, dist};
            P -= D;
            sum01 += P.sum + P.cnt * D.sum;
            P -= D;
        }
    };
    auto reset = [&](u64 v) -> void {
        // B で頂点 v の重みを 0 に
        auto& anc = centroid_ancestor[v];
        T D;
        {
            auto [cen, dist] = anc[0];
            auto& [P, Q] = *cen;
            D = {1, dist};
            P += D;
            sum01 -= P.sum + P.cnt * D.sum;
            P += D;
        }
        for (auto [cen, dist] : anc | views::drop(1)) {
            auto& [P, Q] = *cen;
            Q += D;
            sum01 += Q.sum + Q.cnt * D.sum;
            Q += D;
            D = {1, dist};
            P += D;
            sum01 -= P.sum + P.cnt * D.sum;
            P += D;
        }
    };
    auto set_subtree = [&](auto set_subtree, u64 v) -> void {
        // 部分木の重みをすべて 1 に
        set(v);
        for (u64 i : A[v]) {
            set_subtree(set_subtree, i);
        }
    };
    auto reset_subtree = [&](auto reset_subtree, u64 v) -> void {
        // 部分木の重みをすべて 0 に
        reset(v);
        for (u64 i : A[v]) {
            reset_subtree(reset_subtree, i);
        }
    };
    auto climb = [&](auto climb, u64 v) -> void {
        // 部分木の重みをすべて 1 に
        set(v);
        for (u64 i : A[v] | views::drop(1)) {
            set_subtree(set_subtree, i);
        }
        ans += sum01;
        if (u64 p = heavy_parent[v]; p != -1) {
            // heavy-edge を登る
            climb(climb, p);
        } else {
            reset_subtree(reset_subtree, v);
        }
    };
    auto t6 = steady_clock::now();
    cerr << "pre-calc: " << duration_cast<milliseconds>(t6 - t_Bzero1).count() << "ms\n";
    // heavy-path を下から順に処理
    auto t_calc0 = steady_clock::now();
    for (auto&& [i, a] : A | views::enumerate) {
        if (a.size()) continue;
        climb(climb, (u64)i);
    }
    auto t_calc1 = steady_clock::now();
    cerr << "calc: " << duration_cast<milliseconds>(t_calc1 - t_calc0).count() << "ms\n";
    // 出力
    auto t_out0 = steady_clock::now();
    cout << ans * 2 << '\n';
    auto t_out1 = steady_clock::now();
    cerr << "output: " << duration_cast<milliseconds>(t_out1 - t_out0).count() << "ms\n";
}
int main() {
    pmr::monotonic_buffer_resource mr{ 192 << 20 };
    pmr::set_default_resource(&mr);
    cin.tie(0)->sync_with_stdio(0);
    solve();
}