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

提出ソース
結果

問題	No.3194 Do Optimize Your Solution
ユーザー	👑 tatyam
提出日時	2025-06-27 14:10:06
言語	C++23 (gcc 13.3.0 + boost 1.87.0)
結果	AC
実行時間	951 ms / 3,000 ms
コード長	14,276 bytes
コンパイル時間	5,814 ms
コンパイル使用メモリ	342,976 KB
実行使用メモリ	174,840 KB
最終ジャッジ日時	2025-06-27 21:15:29
合計ジャッジ時間	17,827 ms
ジャッジサーバーID （参考情報）	judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 2
other	AC * 17
権限があれば一括ダウンロードができます
ソースコード

raw source code
#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> eu_in(N), eu_out(N);
    auto t_relabel0 = steady_clock::now();
    {
        u64 idx = 0;
        auto relabel = [&](auto relabel, u64 v) -> void {
            eu_in[v] = idx++;
            for (u64 i : A[v]) {
                relabel(relabel, i);
            }
            eu_out[v] = idx;
        };
        relabel(relabel, 0);
    }
    auto t_relabel1 = steady_clock::now();
    cerr << "A relabel: " << duration_cast<milliseconds>(t_relabel1 - t_relabel0).count() << "ms\n";
    // heavy_parent: A における親の番号 (light-edge なら -1)
    auto t_parent0 = steady_clock::now();
    V<u64> heavy_parent(N, -1);
    V<bool> heavy_leaf(N, 1);
    for (auto&& [i, a] : views::enumerate(A)) {
        if (a.empty() || i == 0) continue; // 根はやらなくていいので切る
        heavy_parent[a[0]] = i;
        heavy_leaf[i] = 0;
    }
    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 = eu_in[get() - 1];
        v = eu_in[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 = std::move(A);
    for (auto&& [b, d] : views::zip(B, deg)) {
        b.clear();
        b.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;
        }
    };
    using Arr = array<T, 2>;
    V<Arr> sum_dist(N); // { (重心が i である連結成分) -> i, (重心が i である連結成分) -> parent[i] }
    struct C { Arr* cen = nullptr; u64 dist = -1; };
    V<V<C>> centroid_ancestor(N); // {sum_dist の index，重心までの距離}
    auto t_centroid0 = steady_clock::now();
    {
        array<V<C>, 19> centroid_ancestor_tp{}; // {sum_dist の index，重心までの距離}
        for (auto& v : centroid_ancestor_tp) v.resize(N);
        V<u64> depth(N, -1); // 重心分解木の深さ
        V<pair<u64, u64>> cc; // {index, 部分木サイズ}
        cc.reserve(N);
        auto adj = [&](u64 v) {
            return B[v] | views::filter([&](u64 u) { return depth[u] == -1; });
        };
        
        auto CD = [&](auto CD, u64 dep, 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;
            // 重心を削除
            depth[v] = dep;
            erase_if(B[v], [&](u64 u) { return depth[u] != -1; });
            // centroid_ancestor を更新しながら，子の連結成分の siz を求める
            const u64 cen = v;
            auto& anc = centroid_ancestor_tp[dep];
            anc[v] = {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 {
                    anc[v] = {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;
                };
                const u64 siz = dfs(dfs, v, u, 1);
                u = CD(CD, dep + 1, u, siz);
            }
            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, 0, N);
        auto t_centroid1 = steady_clock::now();
        cerr << "centroid decomp: " << duration_cast<milliseconds>(t_centroid1 - t_centroid0).count() << "ms\n";
        // 重心分解木 B の relabel
        vector<u64> label(N);
        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&& [v, anc] : centroid_ancestor | views::enumerate) {
            anc.resize(depth[v] + 1);
            for (auto&& [d, x] : views::enumerate(anc)) {
                auto [cen, dist] = centroid_ancestor_tp[d][v];
                x = {&sum_dist[label[bit_cast<u64>(cen)]], dist};
            }
        }
        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;
        }
    };
    pair<u64, u64> subtree_range = {-1, -1};
    auto set_subtree = [&](u64 v) {
        // 部分木の重みをすべて 1 に
        auto& [L, R] = subtree_range;
        if (L == -1) {
            L = eu_in[v];
            R = eu_out[v];
            for (u64 i : views::iota(L, R)) set(i);
            return;
        }
        set(eu_in[v]);
        L = eu_in[v];
        for (u64 i : views::iota(R, eu_out[v])) {
            set(i);
        }
        R = eu_out[v];
    };
    auto reset_subtree = [&] {
        // 部分木の重みをすべて 0 に
        auto [L, R] = subtree_range;
        for (u64 i : views::iota(L, R)) reset(i);
        subtree_range = {-1, -1};
    };
    auto climb = [&](auto climb, u64 v) -> void {
        // 部分木の重みをすべて 1 に
        set_subtree(v);
        ans += sum01;
        if (u64 p = heavy_parent[v]; p != -1) {
            // heavy-edge を登る
            climb(climb, p);
        } else {
            reset_subtree();
        }
    };
    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 (u64 i : views::iota(1u, N)) {
        if (heavy_leaf[i]) 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();
}
yukicoder

結果

ソースコード
