結果

問題 No.3194 Do Optimize Your Solution
ユーザー noya2
提出日時 2025-06-23 23:02:52
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
TLE  
(最新)
AC  
(最初)
実行時間 -
コード長 12,380 bytes
コンパイル時間 3,884 ms
コンパイル使用メモリ 303,452 KB
実行使用メモリ 62,816 KB
最終ジャッジ日時 2025-06-27 20:51:13
合計ジャッジ時間 12,835 ms
ジャッジサーバーID
(参考情報)
judge4 / judge2
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 16 TLE * 1
権限があれば一括ダウンロードができます

ソースコード

diff #

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

#define rep(i,n) for (int i = 0; i < (int)(n); i++)
#define repp(i,m,n) for (int i = (m); i < (int)(n); i++)
#define all(v) (v).begin(),(v).end()
using ull = unsigned long long;

namespace noya2{

/* ~ (. _________ . /) */

}

using namespace noya2;

namespace noya2::internal {

template<class E>
struct csr {
    csr () {}
    csr (int _n) : n(_n) {}
    csr (int _n, int m) : n(_n){
        start.reserve(m);
        elist.reserve(m);
    }
    // ACL style constructor (do not have to call build)
    csr (int _n, const std::vector<std::pair<int,E>> &idx_elem) : n(_n), start(_n + 2), elist(idx_elem.size()) {
        for (auto &[i, e] : idx_elem){
            start[i + 2]++;
        }
        for (int i = 1; i < n; i++){
            start[i + 2] += start[i + 1];
        }
        for (auto &[i, e] : idx_elem){
            elist[start[i + 1]++] = e;
        }
        prepared = true;
    }
    int add(int idx, E elem){
        int eid = start.size();
        start.emplace_back(idx);
        elist.emplace_back(elem);
        return eid;
    }
    void build(){
        if (prepared) return ;
        int m = start.size();
        std::vector<E> nelist(m);
        std::vector<int> nstart(n + 2, 0);
        for (int i = 0; i < m; i++){
            nstart[start[i] + 2]++;
        }
        for (int i = 1; i < n; i++){
            nstart[i + 2] += nstart[i + 1];
        }
        for (int i = 0; i < m; i++){
            nelist[nstart[start[i] + 1]++] = elist[i];
        }
        swap(elist,nelist);
        swap(start,nstart);
        prepared = true;
    }
    const auto operator[](int idx) const {
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    auto operator[](int idx){
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    const auto operator()(int idx, int l, int r) const {
        return std::ranges::subrange(elist.begin()+start[idx]+l,elist.begin()+start[idx]+r);
    }
    auto operator()(int idx, int l, int r){
        return std::ranges::subrange(elist.begin()+start[idx]+l,elist.begin()+start[idx]+r);
    }
    size_t size() const {
        return n;
    }
    int n;
    std::vector<int> start;
    std::vector<E> elist;
    bool prepared = false;
};

} // namespace noya2::internal

struct simple_tree_for_yuki {
    internal::csr<int> g;
    simple_tree_for_yuki (int n) : g(n, (n - 1)*2) {}
    template<typename Stream>
    void input_edges(Stream &inputter){
        for (int i = 0; i < g.n - 1; i++){
            int u, v; inputter >> u >> v;
            u--, v--;
            g.add(u, v);
            g.add(v, u);
        }
        g.build();
    }
    const auto operator[](int v) const {
        return g[v];
    }
    auto operator[](int v){
        return g[v];
    }
    int size() const {
        return g.n;
    }
};

struct hld_tree_for_yuki {
    int n, root;
    std::vector<int> down, nxt, sub, tour, dep;
    hld_tree_for_yuki (int _n) : n(_n), root(0), down(n), nxt(n), sub(n, 1), tour(n), dep(n) {}
    template<typename Stream>
    void input_edges(Stream &inputter){
        for (int i = 1; i < n; i++){
            int u, v; inputter >> u >> v;
            u--, v--;
            down[u]++;
            down[v]++;
            nxt[u] ^= v;
            nxt[v] ^= u;
        }
        build_from_edges();
    }
    ull depth(int v) const {
        return dep[v];
    }
    int leader(int v) const {
        return nxt[v] < 0 ? v : nxt[v];
    }
    // down[v] : degree of v
    // nxt[v] : xor prod of neighbor of v
    void build_from_edges(){
        // use tour as queue
        int back = 0;
        for (int u = 0; u < n; u++){
            if (u != root && down[u] == 1){
                tour[back++] = u;
            }
        }
        for (int front = 0; front < n - 1; front++){
            int u = tour[front];
            down[u] = -1;
            int v = nxt[u]; // parent of v
            nxt[v] ^= u;
            if (--down[v] == 1 && v != root){
                tour[back++] = v;
            }
        }
        // check : now, tour is reverse of topological order

        tour.pop_back();

        // check : now, down[*] <= 1
        for (int u : tour){
            int v = nxt[u];
            // subtree size (initialized (1,1,...,1))
            sub[v] += sub[u];
            // heaviest subtree of its child
            if (down[v] < sub[u]){
                down[v] = sub[u];
            }
        }
        for (int u : tour){
            int v = nxt[u];
            // whether u is not the top of heavy path
            if (down[v] == sub[u]){
                sub[u] = ~sub[u];
                down[v] = ~down[v];
            }
        }

        // after appearing v as u (or v == root), 
        // down[v] is the visiting time of euler tour
        // nxt[v] is the lowest vertex of heavy path which contains v
        //   (if v itself, nxt[v] is ~(parent of v))
        // sub[v] + down[v] is the light child's starting time of euler tour
        // note : heavy child's visiting time of euler tour is (the time of its parent) + 1
        sub[root] = ~down[root] + 1;
        down[root] = 0;
        nxt[root] = -1;
        for (int u : tour | std::views::reverse){
            int v = nxt[u];
            int nsub = ~down[u] + 1;
            // heavy child
            if (sub[u] < 0){
                down[u] = down[v] + 1;
                nxt[u] = (nxt[v] < 0 ? v : nxt[v]);
            }
            // light child
            else {
                down[u] = down[v] + sub[v];
                sub[v] += sub[u];
                nxt[u] = ~v;
            }
            sub[u] = nsub;
            dep[u] = dep[v] + 1;
        }

        // tour is inverse permutation of down
        tour.push_back(root);
        for (int u = 0; u < n; u++){
            tour[down[u]] = u;
        }
    }
};

namespace noya2 {

std::vector<int> centroid_decomposition(const auto &g){
    int n = g.size();
    if (n == 0){
        return {};
    }
    std::vector<int> sub(n), order;
    order.reserve(n);
    auto subtree = [&](auto sfs, int v, int f) -> void {
        sub[v] = 1;
        for (int u : g[v]){
            if (u == f) continue;
            sfs(sfs, u, v);
            sub[v] += sub[u];
        }
    };
    subtree(subtree,0,-1);
    auto fixed_root = [&](auto self, int root, int par, int cur_size) -> void {
        auto dfs = [&](auto sfs, int v, int f, int sz) -> int {
            int heavy = 0, child = -1;
            for (int u : g[v]){
                if (u == f) continue;
                if (heavy < sub[u]){
                    heavy = sub[u];
                    child = u;
                }
            }
            if (heavy > sz/2){
                int ret = sfs(sfs, child, v, sz);
                sub[v] -= ret;
                return ret;
            }
            else {
                order.emplace_back(v);
                for (int u : g[v]){
                    if (u == f) continue;
                    self(self, u, v, sub[u]);
                }
                int ret = sub[v];
                sub[v] = 0;
                return ret;
            }
        };
        while (cur_size > 0){
            cur_size -= dfs(dfs, root, par, cur_size);
        }
    };
    fixed_root(fixed_root, 0, -1, n);
    return order;
}

} // namespace noya2

int main(){
    cin.tie(0)->sync_with_stdio(0);
    int n; cin >> n;
    simple_tree_for_yuki a(n);
    hld_tree_for_yuki b(n);
    a.input_edges(cin);
    b.input_edges(cin);
    vector<int> spfensz(n);
    rep(v,n){
        int lv = b.leader(v);
        spfensz[lv] = max(spfensz[lv], b.down[v] - b.down[lv] + 1);
    }
    vector<array<ull,2>> spfenfront(n,array<ull,2>{});
    vector<array<ull,4>> spfen(n,array<ull,4>{});
    ull rootdepsum = 0, rootdistsum = 0;
    auto apply_from_root = [&](int v, int dist, int sgn){
        ull x = 2 * sgn, y = 2 * dist * sgn;
        rootdepsum += x;
        rootdistsum += y;
        while (true){
            int lv = b.leader(v);
            int start = b.down[lv];
            int sz = spfensz[lv];
            int r = b.down[v] + 1 - start;
            int memor = r;
            r++;
            while (r <= sz){
                int id = r - 1 + start;
                spfen[id][0] += x * memor;
                spfen[id][1] += -x;
                spfen[id][2] += y * memor;
                spfen[id][3] += -y;
                r += r & -r;
            }
            spfenfront[start][0] += x;
            spfenfront[start][1] += y;
            if (lv == 0) break;
            v = ~b.nxt[lv];
        }
    };
    auto prod_from_root = [&](int v, int dist){
        ull depsum = -rootdepsum, distsum = -rootdistsum;
        while (true){
            int lv = b.leader(v);
            int start = b.down[lv];
            int r = b.down[v] + 1 - start;
            int memor = r;
            array<ull,4> prod = {};
            while (r > 0){
                int id = r - 1 + start;
                rep(i,4){
                    prod[i] += spfen[id][i];
                }
                r -= r & -r;
            }
            depsum += prod[0] + (prod[1] + spfenfront[start][0]) * memor;
            distsum += prod[2] + (prod[3] + spfenfront[start][1]) * memor;
            if (lv == 0) break;
            v = ~b.nxt[lv];
        }
        return depsum * dist + distsum;
    };
    vector<bool> done(n,false);
    ull ans = 0;
    for (int ctr : centroid_decomposition(a)){
        done[ctr] = true;
        vector<pair<int,int>> vds = {{ctr,0}};
        vector<int> start = {0,1};
        auto dfs = [&](auto sfs, int v, int f, int dist) -> void {
            vds.emplace_back(v,dist);
            for (int u : a[v]){
                if (done[u]) continue;
                if (u == f) continue;
                sfs(sfs,u,v,dist+1);
            }
        };
        for (int v : a[ctr]){
            if (done[v]) continue;
            dfs(dfs,v,ctr,1);
            start.emplace_back(vds.size());
        }
        // numof u, dist[u], dep[u], dist[u]*dep[u]
        ull cnt = 0, distsum = 0, depsum = 0, prdsum = 0;
        int csz = start.size() - 2;
        // v : ctr
        [&]{
            // dist[u] * dep[u]
            prdsum += 0;
            // dist[u] * dep[v]
            distsum += 0;
            // dist[v] * dep[u]
            depsum += b.depth(ctr);
            // dist[v] * dep[v]
            cnt += 1;
            // dist[u] * -2dep[lca(u,v)]
            // dist[v] * -2dep[lca(u,v)]
            apply_from_root(ctr, 0, -1);
        }();
        auto proc1 = [&](int v, int dist){
            // dist[u] * dep[u]
            ans += prdsum;
            // dist[u] * dep[v]
            ans += distsum * b.depth(v);
            // dist[v] * dep[u]
            ans += dist * depsum;
            // dist[v] * dep[v]
            ans += (ull)dist * b.depth(v) * cnt;
            // dist[u] * -2dep[lca(u,v)]
            // dist[v] * -2dep[lca(u,v)]
            ans += prod_from_root(v, dist);
        };
        auto proc2 = [&](int v, int dist){
            // dist[u] * dep[u]
            prdsum += (ull)dist * b.depth(v);
            // dist[u] * dep[v]
            distsum += dist;
            // dist[v] * dep[u]
            depsum += b.depth(v);
            // dist[v] * dep[v]
            cnt += 1;
            // dist[u] * -2dep[lca(u,v)]
            // dist[v] * -2dep[lca(u,v)]
            apply_from_root(v, dist, -1);
        };
        auto proc3 = [&](int v, int dist){
            // cancel of
            // dist[u] * -2dep[lca(u,v)]
            // dist[v] * -2dep[lca(u,v)]
            apply_from_root(v, dist, 1);
        };
        rep(i,csz){
            repp(j,start[i+1],start[i+2]){
                auto [v, dist] = vds[j];
                proc1(v,dist);
            }
            repp(j,start[i+1],start[i+2]){
                auto [v, dist] = vds[j];
                proc2(v,dist);
            }
        }
        rep(i,csz){
            repp(j,start[i+1],start[i+2]){
                auto [v, dist] = vds[j];
                proc3(v,dist);
            }
        }
        [&]{
            // cancel of
            // dist[u] * -2dep[lca(u,v)]
            // dist[v] * -2dep[lca(u,v)]
            apply_from_root(ctr, 0, 1);
        }();
    }
    ans *= 2;
    cout << ans << endl;
}
0