Submission #4666035


Source Code Expand

Copy
//#undef LOCAL
#include <bits/stdc++.h>
#include <ostream>

using namespace std;
using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
template<class T> using V = vector<T>;
template<class T> using VV = V<V<T>>;
constexpr ll TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n-1); }
#define rep(i,N) for(int i=0;i<(int)(N);i++)
#define rep1(i,N) for(int i=1;i<=(int)(N);i++)
#define fs first
#define sc second
#define pb push_back
#define PB push_back
#define MP make_pair
#define FOR(i, a, b) for(int i=int(a);i<int(b);i++)
#define REP(i,b) FOR(i,0,b)
#define all(x) x.begin(),x.end()
#define ALL(x) x.begin(),x.end()
template<class T, class U> void chmin(T& t, const U& u) { if (t > u) t = u; }
template<class T, class U> void chmax(T& t, const U& u) { if (t < u) t = u; }

#ifdef LOCAL
#define show(x) cerr << __LINE__ << " : " << #x << " = " << (x) << endl
#else
#define show(x) true
#endif

template <class T, class U>
ostream& operator<<(ostream& os, const pair<T, U>& p) {
    return os << "P(" << p.first << ", " << p.second << ")";
}

template <class T> ostream& operator<<(ostream& os, const V<T>& v) {
    os << "[";
    for (auto d : v) os << d << ", ";
    return os << "]";
}

template <class N, class E> struct AllTree {
    int n;
    const VV<E>& g;
    V<N> sm;
    VV<N> dp; // tree
    void dfs1(int p, int b) {
        sm[p] = N();
        for (auto e : g[p]) {
            int d = e.to;
            if (d == b) continue;
            dfs1(d, p);
            sm[p] = sm[p] + sm[d].to_subs(p, e);
        }
        sm[p] = sm[p].join(p);
    }
    void dfs2(int p, int b, N top) {
        int deg = int(g[p].size());
        dp[p] = V<N>(deg + 1);
        dp[p][0] = N();
        for (int i = 0; i < deg; i++) {
            int d = g[p][i].to;
            dp[p][i + 1] = dp[p][i] + (d == b ? top : sm[d]).to_subs(p, g[p][i]);
        }
        N rnode = N();
        dp[p].back() = dp[p].back().join(p);
        for (int i = deg - 1; i >= 0; i--) {
            dp[p][i] = (dp[p][i] + rnode).join(p);
            int d = g[p][i].to;
            if (d != b) dfs2(d, p, dp[p][i]);
            rnode = rnode + (d == b ? top : sm[d]).to_subs(p, g[p][i]);
        }
    }
    AllTree(const VV<E>& _g) : n(int(_g.size())), g(_g), sm(n), dp(n) {
        dfs1(0, -1);
        dfs2(0, -1, N());
    }
};

template <class N, class E> VV<N> get_all_tree(const VV<E>& g) {
    return AllTree<N, E>(g).dp;
}

struct Node {
    // Diameter of Tree
    int rad = 0, dia = 0; // radius(tree), diameter
    array<int, 2> rd = {{0, 0}}; // radiuses(subtrees)
    template <class E> Node to_subs(int, const E& e) const {
        // tree -> subtrees
        return {-1, dia, {rad + e.dist, 0}};
    }
    Node operator+(const Node& r) const {
        // subtrees + subtrees
        array<int, 4> v = {rd[0], rd[1], r.rd[0], r.rd[1]};
        sort(v.begin(), v.end(), greater<>());
        return {-1, max(dia, r.dia), {v[0], v[1]}};
    }
    Node join(int) const {
        // subtrees -> tree
        return {rd[0], max(dia, rd[0] + rd[1]), {}};
    }

    friend ostream &operator<<(ostream &os, const Node &node) {
        os << "rad: " << node.rad << " dia: " << node.dia << " rd: " << node.rd[0] << " " << node.rd[1];
        return os;
    }
};

struct E {
    int to, dist;
};

int main() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << setprecision(20) << fixed;
    int n;
    cin >> n;
    VV<E> ga(n);
    for (int i = 0; i < n-1; i++) {
        int a, b;
        cin >> a >> b; a--; b--;
        ga[a].push_back({b, 1});
        ga[b].push_back({a, 1});
    }
    int m;
    cin >> m;
    VV<E> gb(m);
    for (int i = 0; i < m-1; i++) {
        int a, b;
        cin >> a >> b; a--; b--;
        gb[a].push_back({b, 1});
        gb[b].push_back({a, 1});
    }

    auto tra = get_all_tree<Node>(ga), trb = get_all_tree<Node>(gb);

    V<ll> va, vb;
    ll di = -1;
    for (int i = 0; i < n; i++) {
        va.push_back(tra[i].back().rad);
        di = max<ll>(di, tra[i].back().dia);
    }
    for (int i = 0; i < m; i++) {
        vb.push_back(trb[i].back().rad);
        di = max<ll>(di, trb[i].back().dia);
    }

    show(di);
    show(va);
    show(vb);

    sort(va.begin(), va.end());
    sort(vb.begin(), vb.end());

    V<ll> vbs(m + 1);
    for (int i = 0; i < m; i++) {
        vbs[i + 1] = vbs[i] + vb[i];
    }
    ll ans = 0;
    for (ll d: va) {
        d++;
        //d + x <= di
        int lw = lower_bound(vb.begin(), vb.end(), di - d) - vb.begin();
        ans += ll(lw) * di;
        ans += (vbs.back() - vbs[lw]) + ll(m - lw) * d;
    }
    cout << ans << endl;
    return 0;
}

Submission Info

Submission Time
Task B - Bonsai Grafting
User yosupo
Language C++14 (GCC 5.4.1)
Score 700
Code Size 4825 Byte
Status
Exec Time 190 ms
Memory 42340 KB

Judge Result

Set Name Score / Max Score Test Cases
Sample 0 / 0 01.txt, 02.txt
All 700 / 700 01.txt, 02.txt, 11.txt, 12.txt, 13.txt, 14.txt, 15.txt, 16.txt, 17.txt, 18.txt, 19.txt, 20.txt, 21.txt, 22.txt, 23.txt, 24.txt, 25.txt, 26.txt, 27.txt, 28.txt, 29.txt, 30.txt, 31.txt, 32.txt, 33.txt, 34.txt, 35.txt, 36.txt, 37.txt, 38.txt, 39.txt, 40.txt, 41.txt, 42.txt, 43.txt, 44.txt, 45.txt, 46.txt, 47.txt, 48.txt, 49.txt, 50.txt
Case Name Status Exec Time Memory
01.txt 1 ms 256 KB
02.txt 1 ms 256 KB
11.txt 1 ms 256 KB
12.txt 1 ms 256 KB
13.txt 1 ms 256 KB
14.txt 1 ms 256 KB
15.txt 1 ms 256 KB
16.txt 1 ms 256 KB
17.txt 1 ms 256 KB
18.txt 1 ms 256 KB
19.txt 1 ms 256 KB
20.txt 1 ms 256 KB
21.txt 164 ms 32740 KB
22.txt 162 ms 32868 KB
23.txt 181 ms 32612 KB
24.txt 187 ms 40292 KB
25.txt 158 ms 32868 KB
26.txt 169 ms 32740 KB
27.txt 181 ms 40164 KB
28.txt 169 ms 32484 KB
29.txt 179 ms 40292 KB
30.txt 190 ms 41828 KB
31.txt 165 ms 32740 KB
32.txt 161 ms 32868 KB
33.txt 169 ms 32612 KB
34.txt 175 ms 42340 KB
35.txt 155 ms 32868 KB
36.txt 176 ms 32740 KB
37.txt 182 ms 37860 KB
38.txt 162 ms 32484 KB
39.txt 175 ms 39268 KB
40.txt 185 ms 39652 KB
41.txt 78 ms 16740 KB
42.txt 78 ms 16740 KB
43.txt 73 ms 16740 KB
44.txt 78 ms 16740 KB
45.txt 75 ms 16868 KB
46.txt 70 ms 16868 KB
47.txt 91 ms 26084 KB
48.txt 80 ms 16740 KB
49.txt 88 ms 21860 KB
50.txt 90 ms 21860 KB