結果

問題 No.3194 Do Optimize Your Solution
ユーザー 👑 tatyam
提出日時 2025-06-24 17:56:00
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
MLE  
実行時間 -
コード長 6,570 bytes
コンパイル時間 5,142 ms
コンパイル使用メモリ 339,500 KB
実行使用メモリ 630,112 KB
最終ジャッジ日時 2025-06-27 20:51:37
合計ジャッジ時間 23,885 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 9 MLE * 1 -- * 7
権限があれば一括ダウンロードができます

ソースコード

diff # 

#ifdef DEBUG
#include <range/v3/all.hpp>
#endif
#include <bits/stdc++.h>
using namespace std;
using u32 = uint32_t;
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 を下から順に，
 */
void solve() {
    #ifdef DEBUG
    namespace views = ::ranges::views;
    namespace ranges = std::ranges;
    #endif
    // 入力
    u64 N;
    cin >> N;
    // 実質的に CSR な読み込み
    vector<pair<u32, u32>> edge(N - 1);
    vector<u32> deg(N);
    for (auto& [u, v] : edge) {
        cin >> u >> v;
        u--; v--;
        deg[u]++;
        deg[v]++;
    }
    vector<V<u32>> 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);
    }
    deg.assign(N, 0);
    for (auto& [u, v] : edge) {
        cin >> u >> v;
        u--; v--;
        deg[u]++;
        deg[v]++;
    }
    vector<V<u32>> 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);
    }
    
    // A を HL 分解
    // heavy-edge は G[*][0] に
    vector<u32> siz(N, 1), heavy_parent(N, -1);
    auto hld = [&](auto dfs, u32 v) -> void {
        if (A[v].empty()) return;
        for (u32 i : A[v]) {
            erase(A[i], v);
            dfs(dfs, i);
            siz[v] += siz[i];
        }
        auto p = ranges::max_element(A[v], {}, [&](u32 i) { return siz[i]; });
        if (v) heavy_parent[*p] = v;  // 根はやらなくて良い
        swap(*p, A[v][0]);
    };
    hld(hld, 0);
    // B を重心分解
    struct T { u64 cnt = 0, sum = 0; };
    vector<V<T>> sum_dist(N);
    for (auto&& [b, s] : views::zip(B, sum_dist)) s.resize(b.size() + 1, {});
    vector<V<tuple<T*, T*, u64>>> centroid_ancestor(N); // {連結成分の和，重心からその方向の和，重心までの距離}
    for (auto& c : centroid_ancestor) c.reserve(bit_width(N));
    {
        vector<bool> deleted(N, false);
        vector<pair<u32, u32>> cc; // {index, 部分木サイズ}
        auto adj = [&](u32 v) {
            return B[v] | views::filter([&](u32 u) { return !deleted[u]; });
        };
        auto dfs_size = [&](auto dfs_size, u32 parent, u32 v) -> u32 {
            // 頂点 v の部分木のサイズを求めながら，頂点のリストを作る
            u32 siz = 1;
            for (u32 i : adj(v)) if (i != parent) {
                siz += dfs_size(dfs_size, v, i);
            }
            cc.emplace_back(v, siz);
            return siz;
        };

        auto CD = [&](auto CD, u32 v) -> void {
            // 頂点 v を含む連結成分を重心分解する
            cc.clear();
            u32 half_siz = dfs_size(dfs_size, -1, v) / 2;
            const u32 cen = ranges::max(cc, {}, [&](const auto& p) { return p.second <= half_siz ? p.second : 0; }).first;
            // centroid_ancestor を更新
            auto& S = sum_dist[cen];
            T* const p_sum = &S.back();
            centroid_ancestor[cen].emplace_back(p_sum, nullptr, 0);
            for (auto&& [i, s] : views::zip(B[cen], S)) {
                if (deleted[i]) continue;
                T* const q_sum = &s;
                auto dfs = [&](auto dfs, u32 parent, u32 v, u32 dist) -> void {
                    centroid_ancestor[v].emplace_back(p_sum, q_sum, dist);
                    for (u32 u : adj(v)) if  (u != parent) {
                        dfs(dfs, v, u, dist + 1);
                    }
                };
                dfs(dfs, cen, i, 1);
            }
            // 重心を削除
            deleted[cen] = true;
            // 子を重心分解
            for (u32 u : adj(cen)) CD(CD, u);
        };
        CD(CD, 0);
    }
    // B の全頂点に 0 をセット
    for (auto& anc : centroid_ancestor) {
        for (auto& [p_sum, q_sum, dist] : anc) {
            auto& P = *p_sum;
            P.cnt++;
            P.sum += dist;
            if (q_sum == nullptr) break;
            auto& Q = *q_sum;
            Q.cnt++;
            Q.sum += dist;
        }
    }
    vector<bool> val(N);
    u64 ans = 0, sum01 = 0;
    auto add = [&](u32 v, int x) -> void {
        if (val[v]) assert(x == -1);
        else assert(x == 1);
        val[v] = !val[v];
        // B で頂点 v の重みを 1 に / 0 に
        for (auto& [p_sum, q_sum, dist] : centroid_ancestor[v]) {
            auto& P = *p_sum;
            P.cnt -= x;
            P.sum -= x * dist;
            sum01 += x * (P.sum + P.cnt * dist);
            P.cnt -= x;
            P.sum -= x * dist;
            if (q_sum == nullptr) break;
            auto& Q = *q_sum;
            Q.cnt -= x;
            Q.sum -= x * dist;
            sum01 -= x * (Q.sum + Q.cnt * dist);
            Q.cnt -= x;
            Q.sum -= x * dist;
        }
    };
    auto set_subtree = [&](auto set_subtree, u32 v) -> void {
        // 部分木の重みをすべて 1 に
        add(v, 1);
        for (u32 i : A[v]) {
            set_subtree(set_subtree, i);
        }
    };
    auto reset_subtree = [&](auto reset_subtree, u32 v) -> void {
        // 部分木の重みをすべて 0 に
        add(v, -1);
        for (u32 i : A[v]) {
            reset_subtree(reset_subtree, i);
        }
    };
    auto climb = [&](auto climb, u32 v) -> void {
        // 部分木の重みをすべて 1 に
        add(v, 1);
        for (u32 i : A[v] | views::drop(1)) {
            set_subtree(set_subtree, i);
        }    
        ans += sum01;
        if (u32 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, (u32)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();
}
0