2016北京集训测试赛（八）Problem C: 直径

Solution

一个定理: 把两棵树用一条边练成一棵树后, 树的直径在原来两棵树的四个直径端点中产生.

放到这一题, 我们通过DP先求出大树中以每个点为根的子树中的直径, 再取每棵小树中与其他树有连边的点以及两个直径端点作为虚树上的关键点, 建虚树再求一次直径即可.

#include <cstdio>
#include <cctype>
#include <vector>
#include <map>
#include <algorithm>
#define vector std::vector
#define map std::map
#define sort std::sort
#define swap std::swap

namespace Zeonfai
{
    inline int getInt()
    {
        int a = 0, sgn = 1;
        char c;
        while(! isdigit(c = getchar())) if(c == '-') sgn *= -1;
        while(isdigit(c)) a = a * 10 + c - '0', c = getchar();
        return a * sgn;
    }
}
const int N = (int)3e5, M = (int)3e5, LOG = 19;
vector<int> ky[M + 1];
struct virtualTree
{
    struct node;
    struct edge
    {
        node *v; int w;
        inline edge(node *_v, int _w) {v = _v; w = _w;}
    };
    struct node
    {
        vector<edge> edg;
        long long dep, maxLen;
        node *dmt[2];
        inline node() {edg.clear();}
    } *rt;
    inline void addEdge(node *u, node *v, int len)
    {
        u->edg.push_back(edge(v, len)); v->edg.push_back(edge(u, len));
    }
    node* DFS(node *u, node *pre)
    {
        u->dmt[0] = u->dmt[1] = u;
        for(auto edg : u->edg) if(edg.v != pre)
        {
            node *v = edg.v; v->dep = u->dep + edg.w;
            node *res = DFS(v, u);
            if(res->dep > u->dmt[0]->dep) u->dmt[1] = u->dmt[0], u->dmt[0] = res; else if(res->dep > u->dmt[1]->dep) u->dmt[1] = res;
        }
        u->maxLen = u->dmt[0]->dep + u->dmt[1]->dep - (u->dep << 1);
        node *res = u->dmt[0];
        for(auto edg : u->edg) if(edg.v != pre && edg.v->maxLen > u->maxLen) u->dmt[0] = edg.v->dmt[0], u->dmt[1] = edg.v->dmt[1], u->maxLen = edg.v->maxLen;
        return res;
    }
    inline long long getAnswer()
    {
        rt->dep = 0;
        DFS(rt, NULL);
        return rt->maxLen;
    }
}vir;
struct edge
{
    int id[2], u[2];
}edg[M];
struct node
{
    vector<node*> edg;
    int dep, dfn, maxLen;
    node *dmt[2], *anc[LOG];
    map<int, virtualTree::node*> mp;
    inline node() {edg.clear(); mp.clear();}
};
inline int cmp(node *a, node *b) {return a->dfn < b->dfn;}
struct tree
{
    node nd[N + 1];
    inline void addEdge(int u, int v) {nd[u].edg.push_back(nd + v); nd[v].edg.push_back(nd + u);}
    int clk;
    node* DFS(node *u, node *pre)
    {
        u->dfn = clk ++; u->dep = pre == NULL ? 0 : pre->dep + 1;
        u->dmt[0] = u->dmt[1] = u;
        u->anc[0] = pre;
        for(int i = 1; i < LOG; ++ i) u->anc[i] = u->anc[i - 1] == NULL ? NULL : u->anc[i - 1]->anc[i - 1];
        for(auto v : u->edg) if(v != pre)
        {
            node *res = DFS(v, u);
            if(res->dep > u->dmt[0]->dep) u->dmt[1] = u->dmt[0], u->dmt[0] = res; else if(res->dep > u->dmt[1]->dep) u->dmt[1] = res;
        }
        u->maxLen = u->dmt[0]->dep + u->dmt[1]->dep - (u->dep << 1);
        node *res = u->dmt[0];
        for(auto v : u->edg) if(v != pre && v->maxLen > u->maxLen) u->dmt[0] = v->dmt[0], u->dmt[1] = v->dmt[1], u->maxLen = v->maxLen;
        return res;
    }
    inline void pretreat()
    {
        clk = 0;
        DFS(nd + 1, NULL);
    }
    inline node* getLCA(node *u, node *v)
    {
        if(u->dep < v->dep) swap(u, v);
        for(int i = LOG - 1; ~ i; -- i) if(u->dep - (1 << i) >= v->dep) u = u->anc[i];
        if(u == v) return u;
        for(int i = LOG - 1; ~ i; -- i) if(u->anc[i] != v->anc[i]) u = u->anc[i], v = v->anc[i];
        return u->anc[0];
    }
    inline void buildVirtualTree(int id)
    {
        static node* a[N]; int cnt = 0;
        for(auto u : ky[id]) a[cnt ++] = nd + u;
        sort(a, a + cnt, cmp);
        static node *stk[N]; int tp = 0;
        for(int i = 0; i < cnt; ++ i) if(! i || a[i] != a[i - 1])
        {
            node *u = a[i]; vir.rt = u->mp[id] = new virtualTree::node;
            if(tp > 1)
            {
                node *LCA = getLCA(u, stk[tp - 1]);
                if(LCA != stk[tp - 1])
                {
                    while(tp >= 2 && stk[tp - 2]->dep >= LCA->dep) vir.addEdge(stk[tp - 2]->mp[id], stk[tp - 1]->mp[id], stk[tp - 1]->dep - stk[tp - 2]->dep), -- tp;
                    if(stk[tp - 1] != LCA)
                    {
                         LCA->mp[id] = new virtualTree::node;
                        vir.addEdge(LCA->mp[id], stk[tp - 1]->mp[id], stk[tp - 1]->dep - LCA->dep);
                        -- tp;
                        stk[tp ++] = LCA;
                    }
                }
            }
            stk[tp ++] = u;
        }
        for(int i = 0; i < tp - 1; ++ i) vir.addEdge(stk[i]->mp[id], stk[i + 1]->mp[id], stk[i + 1]->dep - stk[i]->dep);
    }
    inline void addVirtualTreeEdge(int id)
    {
        vir.addEdge(nd[edg[id].u[0]].mp[edg[id].id[0]], nd[edg[id].u[1]].mp[edg[id].id[1]], 1);
    }
}T;
int main()
{

    #ifndef ONLINE_JUDGE

    freopen("diameter.in", "r", stdin);
    freopen("diameter.out", "w", stdout);

    #endif

    using namespace Zeonfai;
    int n = getInt(), m = getInt();
    for(int i = 1; i < n; ++ i)
    {
        int u = getInt(), v = getInt();
        T.addEdge(u, v);
    }
    T.pretreat();
    static int rt[M + 1];
    for(int i = 1; i <= m; ++ i) rt[i] = getInt(), ky[i].push_back(rt[i]), ky[i].push_back(T.nd[rt[i]].dmt[0] - T.nd), ky[i].push_back(T.nd[rt[i]].dmt[1] - T.nd);
    for(int i = 1; i < m; ++ i) for(int j = 0; j < 2; ++ j) edg[i].id[j] = getInt(), edg[i].u[j] = getInt(), ky[edg[i].id[j]].push_back(edg[i].u[j]);
    for(int i = 1; i <= m; ++ i) T.buildVirtualTree(i);
    for(int i = 1; i < m; ++ i) T.addVirtualTreeEdge(i);
    printf("%lld\n", vir.getAnswer() + 1);
}