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

• 問題: No.3194 Do Optimize Your Solution
• ユーザー: 👑 tatyam
• 提出日時: 2025-06-26 08:17:04
• 言語: C++23 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 1,127 ms / 3,000 ms
• コード長: 7,962 bytes
• コンパイル時間: 4,931 ms
• コンパイル使用メモリ: 335,536 KB
• 実行使用メモリ: 146,004 KB
• 最終ジャッジ日時: 2025-06-27 20:55:39
• 合計ジャッジ時間: 19,790 ms
• ジャッジサーバーID （参考情報）: judge3 / judge1

ファイルパターン	結果
sample	AC * 2
other	AC * 17

ソースコード

#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() {
    #ifdef DEBUG
    namespace views = ::ranges::views;
    namespace ranges = std::ranges;
    #endif
    // 入力
    u64 N = get();
    // 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]++;
    }
    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);
    }
    // A を HL 分解
    // heavy-edge は G[*][0] に
    V<u64> siz(N, 1);
    siz.assign(N, 1);
    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);
    // A のオイラーツアー順に頂点番号の relabel
    V<u64> label(N, 0);
    {
        u64 idx = 0;
        auto relabel = [&](auto relabel, u64 v) -> void {
            label[v] = idx++;
            for (u64 i : A[v]) {
                relabel(relabel, i);
            }
        };
        relabel(relabel, 0);
    }
    // A の作り直し
    {
        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_);
    }
    // heavy_parent: A における親の番号 (light-edge なら -1)
    V<u64> heavy_parent(N, -1);
    for (auto&& [i, a] : views::enumerate(A)) {
        if (a.empty() || i == 0) continue; // 根はやらなくていいので切る
        heavy_parent[a[0]] = i;
    }
    // B を入力
    deg.assign(N, 0);
    for (auto& [u, v] : edge) {
        u = label[get() - 1];
        v = label[get() - 1];
        deg[u]++;
        deg[v]++;
    }
    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);
    }
    // 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));
    {
        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 dfs_size = [&](auto dfs_size, u64 parent, u64 v) -> u64 {
            // 頂点 v の部分木のサイズを求めながら，頂点のリストを作る
            u64 siz = 1;
            for (u64 i : adj(v)) if (i != parent) {
                siz += dfs_size(dfs_size, v, i);
            }
            cc.emplace_back(v, siz);
            return siz;
        };

        auto CD = [&](auto CD, u64 v) -> void {
            // 頂点 v を含む連結成分を重心分解する
            cc.clear();
            u64 half_siz = dfs_size(dfs_size, -1, v) / 2;
            v = ranges::min(cc, {}, [&](const auto& p) { return p.second >= half_siz ? p.second : u64(-1); }).first;
            // 重心を削除
            deleted[v] = true;
            // centroid_ancestor を更新
            const auto cen = &sum_dist[v];
            auto dfs = [&](auto dfs, u64 parent, u64 v, u64 dist) -> void {
                centroid_ancestor[v].emplace_back(cen, dist);
                for (u64 u : adj(v)) if  (u != parent) {
                    dfs(dfs, v, u, dist + 1);
                }
            };
            dfs(dfs, -1, v, 0);
            // 子を重心分解
            for (u64 u : adj(v)) CD(CD, u);
        };
        CD(CD, 0);
    }
    // B の全頂点に 0 をセット
    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;
        }
    }
    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);
        }
    };
    // heavy-path を下から順に処理
    for (auto&& [i, a] : A | views::enumerate) {
        if (a.size()) continue;
        climb(climb, (u64)i);
    }
    // 出力
    cout << ans * 2 << '\n';
}
int main() {
    pmr::monotonic_buffer_resource mr{ 256 << 20 };
    pmr::set_default_resource(&mr);
    cin.tie(0)->sync_with_stdio(0);
    solve();
}