Submission #64791673

---

Source Code Expand

#include<bits/stdc++.h>

using namespace std;

typedef vector<int> VI;
typedef long long i64;

int main() {
    cin.tie(nullptr) -> sync_with_stdio(false);
    // freopen("x.in", "r", stdin);
    // freopen("x.out", "w", stdout);

    int n1, n2;
    cin >> n1;
    vector<VI> G1(n1 + 1);
    for(int i = 1, u, v; i < n1; i++) {
        cin >> u >> v;
        G1[u].push_back(v), G1[v].push_back(u);
    }
    cin >> n2;
    vector<VI> G2(n2 + 1);
    for(int i = 1, u, v; i < n2; i++) {
        cin >> u >> v;
        G2[u].push_back(v), G2[v].push_back(u);
    }

    int d1 = 0, d2 = 0;
    VI h1(n1 + 1), h2(n2 + 1);
    auto prework = [&](vector<VI> &G, int &d, VI &h) {
        // 计算直径和树的高度 h
        int s;
        function<void(int, int, int)> dfs = [&](int u, int fa, int dis) {
            if(dis > d) {
                d = dis, s = u;
            }
            for(int v: G[u]) {
                if(v != fa) {
                    dfs(v, u, dis + 1);
                }
            }
        };

        dfs(1, 0, 0), d = 0, dfs(s, 0, 0);

        function<void(int, int)> dp = [&](int u, int fa) {
            for(int v: G[u]) {
                if(v == fa) continue;
                dp(v, u);
                h[u] = max(h[u], h[v] + 1);
            }
        };
        dp(1, 0);
    };

    prework(G1, d1, h1);
    prework(G2, d2, h2);

    // cerr << "d1 = " << d1 << ", d2 = " << d2 << '\n';
    // for(int i = 1; i <= n1; i++) {
        // cerr << h1[i] << " \n"[i == n1];
    // }
    // for(int i = 1; i <= n2; i++) {
        // cerr << h2[i] << " \n"[i == n2];
    // }


    VI f(n1 + 1), g(n2 + 1);
    // f[u]: T1 中离 u 最远的点到 u 的距离
    // g[u]: T2 中离 u 最远的点到 u 的距离
    
    auto calc = [&](vector<VI> &G, VI &ht, VI &dis) {
        auto del = [&](int u, int v) {
            ht[u] = 0;
            for(int _: G[u]) {
                if(_ == v) continue;
                ht[u] = max(ht[u], ht[_] + 1);
            }
        };

        auto add = [&](int u, int v) {
            ht[u] = max(ht[u], ht[v] + 1);
            // cerr << "after add: " << u << ' ' << v << ": ht = " << ht[u] << '\n';
        };

        function<void(int, int)> dp2 = [&](int u, int fa) {
            dis[u] = ht[u];
            // cerr << "In dp2: dis[" << u << "] = " << dis[u] << '\n';
            multiset<int> sonh;
            for(int v: G[u]) {
                sonh.insert(ht[v]);
            }
            for(int v: G[u]) {
                if(v == fa) continue;
                // del(u, v);
                sonh.erase(sonh.find(ht[v]));
                int old = ht[v];
                if(sonh.empty()) {
                    ht[u] = 0;
                } else {
                    ht[u] = 1 + *sonh.rbegin();
                }
                // cerr << "v = " << v << ", ht[u] = " << ht[u] << '\n';
                add(v, u);

                dp2(v, u);
                // del(v, u);
                ht[v] = old;
                sonh.insert(old);
                // add(u, v);
                if(sonh.empty()) {
                    ht[u] = 0;
                } else {
                    ht[u] = 1 + *sonh.rbegin();
                }
            }
        };

        dp2(1, 0);

        // for(int i = 1; i < (int)dis.size(); i++) {
        //     cerr << dis[i] << ' ';
        // }
        // cerr << '\n';
    };

    calc(G1, h1, f), calc(G2, h2, g);

    sort(g.begin() + 1, g.end());

    // cerr << "OK\n";
    vector<i64> sum(n2 + 1);
    for(int i = 1; i <= n2; i++) {
        sum[i] = g[i] + sum[i - 1];
    }
    // for(i64 x: sum) { cerr << x << " "; } cerr << '\n';

    int d = max(d1, d2);
    i64 ans = 0;
    for(int i = 1; i <= n1; i++) {
        int x = d - (1 + f[i]);
        // g[v] > x 时，直径跨越两棵树
        auto it = upper_bound(g.begin() + 1, g.end(), x);
        int p = (int)(it - g.begin());
        i64 add = 1LL * (p - 1) * d;
        if(p != n2 + 1) {
            add += (sum[n2] - sum[p - 1]) + 1LL * (n2 - p + 1) * (f[i] + 1);
        }
        ans += add;
        // cerr << i << ' ' << p << ' ' << add << '\n';
    }

    cout << ans << '\n';
    
    return 0;
}

Submission Info

Compile Error

Main.cpp: In lambda function:
Main.cpp:72:14: warning: variable ‘del’ set but not used [-Wunused-but-set-variable]
   72 |         auto del = [&](int u, int v) {
      |              ^~~

Judge Result