Submission #4665396


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;
    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) {
        show(p);
        show(b);
        show(top);
        int deg = int(g[p].size());
        dp[p] = V<N>(deg + 1);
        dp[p][0] = N();
        show(sm);
        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]);
        }
        show(dp[p]);
        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 {
    // 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().join(-1).rad);
        di = max<ll>(di, tra[i].back().join(-1).dia);
    }
    for (int i = 0; i < m; i++) {
        vb.push_back(trb[i].back().join(-1).rad);
        di = max<ll>(di, trb[i].back().join(-1).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 0
Code Size 4944 Byte
Status

Compile Error

./Main.cpp:94:29: error: array must be initialized with a brace-enclosed initializer
     array<int, 2> rd = {0, 0}; // radiuses(subtrees)
                             ^
./Main.cpp:94:29: error: too many initializers for ‘std::array<int, 2ul>’