Submission #4353051


Source Code Expand

Copy
#include <bits/stdc++.h>

using namespace std;
using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
constexpr ll TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;

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;
    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].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();
        for (int i = deg - 1; i >= 0; i--) {
            int d = g[p][i].to;
            if (d != b) {
                dfs2(d, p, (dp[p][i] + rnode).join(p));
            }
            dp[p][i] = dp[p][i] + rnode;
            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 {
    // サンプル 木の直径
    // firstは 根付き木->森
    // finishは森->根付き木
    int rad, dia;
    int rd[2];
    Node() {
        // !!!!!森を生成すること!!!!!
        rad = 0;
        dia = 0;
        rd[0] = rd[1] = 0;
    }
    template <class E> Node to_subs(int, const E& e) {
        // 根付き木->森 上に辺を追加するイメージ e=(par -> node)
        rd[0] = rad + e.dist;
        rd[1] = 0;
        return *this;
    }
    Node operator+(const Node& r) {
        // 結合則が必須 上に辺を追加された木をまとめるイメージ
        Node n;
        vector<int> v;
        v.push_back(rd[0]);
        v.push_back(rd[1]);
        v.push_back(r.rd[0]);
        v.push_back(r.rd[1]);
        sort(begin(v), end(v), greater<>());
        n.rd[0] = v[0];
        n.rd[1] = v[1];
        n.dia = max(dia, r.dia);
        return n;
    }
    Node join(int) {
        // 元締めとしてidxを追加するイメージ
        rad = rd[0];
        dia = max(dia, rd[0] + rd[1]);
        return *this;
    }
};

struct E {
    int to;
    ll 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().rd[0]);
        di = max<ll>(di, tra[i].back().dia);
    }
    for (int i = 0; i < m; i++) {
        vb.push_back(trb[i].back().rd[0]);
        di = max<ll>(di, trb[i].back().dia);
    }

    /*cout << di << endl;
    cout << va << endl;
    cout << vb << endl;*/

    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 4289 Byte
Status
Exec Time 464 ms
Memory 51300 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 388 ms 35556 KB
22.txt 368 ms 35172 KB
23.txt 366 ms 35684 KB
24.txt 405 ms 48100 KB
25.txt 379 ms 34788 KB
26.txt 387 ms 35428 KB
27.txt 399 ms 47460 KB
28.txt 385 ms 35684 KB
29.txt 439 ms 48228 KB
30.txt 464 ms 51044 KB
31.txt 406 ms 35556 KB
32.txt 405 ms 35172 KB
33.txt 405 ms 35684 KB
34.txt 407 ms 51300 KB
35.txt 367 ms 34788 KB
36.txt 392 ms 35428 KB
37.txt 405 ms 43876 KB
38.txt 389 ms 35684 KB
39.txt 407 ms 46692 KB
40.txt 422 ms 47844 KB
41.txt 183 ms 18148 KB
42.txt 209 ms 18148 KB
43.txt 202 ms 18276 KB
44.txt 196 ms 18148 KB
45.txt 169 ms 17764 KB
46.txt 176 ms 17764 KB
47.txt 212 ms 33380 KB
48.txt 201 ms 18404 KB
49.txt 215 ms 26852 KB
50.txt 241 ms 26852 KB