一道虚树题

题目大意

给树上每个点规定一种颜色

问同色距离最大的点对距离平方是多少

虚树

每次给定树上一个点集(大小为$k$)，你需要构造只包含这些点的树（树的结构必须和原来一样，但是要删去不必要的点），然后回答一些问题。这类题目通常有个特点：$\sum k$规模较小(较小是说能承受$O(\sum klogk)$的复杂度)

对于这题来说，每次相同颜色的点集

对于一个点集，我按照$dfs$序排序，然后一个点一个点往虚树上添加

我维护一个栈(其实就是模拟$dfs$)

当加入一个点$u$时，面临的一种大致情况应该是这样：

当前栈顶元素是$x$，假设弹掉$x$之后栈顶为$y$

我需要不停的弹栈，直到$lca(u,y)=y$，弹栈的过程中不停的连边$(x,y)$

那么这个时候我就要在虚树上添加$lcd(u,x)$这个点，添加完了之后连边$(x,lca(u,x))$，然后$x$也是要弹掉

然后再把$u$入栈

其实上述算法就是模拟了$dfs$的递归和回溯过程

虚树的时间复杂度

对一个大小为$k$的点集建立虚树，可以知道我所添加的$lca$的个数不会超过$k$（可以类比线段树合并，一个$lca$可以看作两个点合并，每次合并会减少一个点，因此合并的次数不超过点数）

所以复杂度是$O(\sum klogk)$的，带$log$是因为我要倍增找$lca$

题解

这题当然可以点分治做

但是正好拿来练习虚树

每次建立出虚树之后，问题就成了找树的直径，一次$dfs$就$ok$了

另外注意这题有坑：存在所有点集大小都等于$1$的数据

代码

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define iinf 0x3f3f3f3f
#define linf (1ll<<60)
#define eps 1e-8
#define maxn 100010
#define maxe 200010
#define maxk 17
#define cl(x) memset(x,0,sizeof(x))
#define rep(i,a,b) for(i=a;i<=b;i++)
#define drep(i,a,b) for(i=a;i>=b;i--)
#define em(x) emplace(x)
#define emb(x) emplace_back(x)
#define emf(x) emplace_front(x)
#define fi first
#define se second
#define de(x) cerr<<#x<<" = "<<x<<endl
using namespace std;
using namespace __gnu_pbds;
typedef long long ll;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
ll read(ll x=0)
{
    ll c, f(1);
    for(c=getchar();!isdigit(c);c=getchar())if(c=='-')f=-f;
    for(;isdigit(c);c=getchar())x=x*10+c-0x30;
    return f*x;
}
struct Graph
{
    int etot, head[maxn], to[maxe], next[maxe], w[maxe];
    void clear(int N)
    {
        for(int i=1;i<=N;i++)head[i]=0;
        etot=0;
    }
    void adde(int a, int b, int c=0){to[++etot]=b;w[etot]=c;next[etot]=head[a];head[a]=etot;}
    #define forp(_,__) for(auto p=__.head[_];p;p=__.next[p])
}G, T;
struct Doubling_LCA
{
    int f[maxn][maxk+1], depth[maxn];
    void clear(int n){for(int i=1;i<=n;i++)depth[i]=0, cl(f[i]);}
    void dfs(Graph &G, int pos, int pre)
    {
        for(auto k=1;(1<<k)<=depth[pos];k++)f[pos][k]=f[f[pos][k-1]][k-1];
        for(auto p(G.head[pos]);p;p=G.next[p])
            if(G.to[p]!=pre)
            {
                f[G.to[p]][0]=pos;
                depth[G.to[p]]=depth[pos]+1;
                dfs(G,G.to[p],pos);
            }
    }
    void run(Graph &G, int root)
    {
        depth[root]=1;
        dfs(G,root,0);
    }
    int q(int x, int y)
    {
        if(depth[x]<depth[y])swap(x,y);
        for(auto k(maxk);~k;k--)
            if(depth[f[x][k]]>=depth[y])
                x=f[x][k];
        if(x==y)return x;
        for(auto k(maxk);~k;k--)
            if(f[x][k]!=f[y][k])
                x=f[x][k], y=f[y][k];
        return f[x][0];
    }
    int jp(int x, int b)
    {
        for(auto k=0;k<=maxk;k++)
            if(b&(1<<k))x=f[x][k];
        return x;
    }
}db;
struct Easy_Tree
{
    int depth[maxn], dist[maxn], tid[maxn], rtid[maxn], tim, size[maxn], rev[maxn];
    void dfs(int pos, int pre, Graph& G)
    {
        tid[pos]=++tim;
        rev[tid[pos]]=pos;
        size[pos]=1;
        forp(pos,G)if(G.to[p]!=pre)
        {
            depth[G.to[p]]=depth[pos]+1;
            dist[G.to[p]]=dist[pos]+G.w[p];
            dfs(G.to[p],pos,G);
            size[pos]+=size[G.to[p]];
        }
        rtid[pos]=tim;
    }
    void run(Graph& G, int root)
    {
        tim=0;
        depth[root]=1;
        dfs(1,0,G);
    }
}et, et2;
struct ConciseTree
{
    //调用build之前，要调用LCA.run()，et.run()
    static ll build(vector<ll>& lis, Graph& T, Doubling_LCA& LCA, Easy_Tree& et)
    {
        sort(lis.begin(),lis.end(),[&](ll a, ll b){return et.tid[a]<et.tid[b];});
        stack<ll> stk;
        auto adde = [&](ll u, ll fa){auto l=et.depth[u]-et.depth[fa];T.adde(u,fa,l);T.adde(fa,u,l);};
        for(auto u:lis)
        {
            if(stk.empty()){stk.em(u);continue;}
            ll lca = LCA.q(u,stk.top());
            while(lca!=stk.top())
            {
                auto x = stk.top(); stk.pop();
                if(stk.empty() or LCA.q(stk.top(),u)==stk.top() and stk.top()!=lca)stk.em(lca);
                adde(x,stk.top());
                lca = LCA.q(u,stk.top());
            }
            stk.em(u);
        }
        while(stk.size()>1)
        {
            auto x=stk.top(); stk.pop();
            adde(x,stk.top());
        }
        return stk.top();
    }
};
ll ans;
ll dfs(Graph& G, ll u, ll fa)
{
    ll mx=0;
    forp(u,G)
    {
        ll v = G.to[p]; if(v==fa)continue;
        ll d = dfs(G,v,u) + G.w[p];
        ans = max( ans, mx+d );
        mx = max(mx,d);
    }
    return mx;
}
ll n, a[maxn];
vector<ll> lis[maxn];
int main()
{
    ll u, v, i, j;
    n=read();
    rep(i,1,n)a[i]=read();
    rep(i,1,n-1)
    {
        u=read(), v=read();
        G.adde(u,v), G.adde(v,u);
    }
    db.run(G,1);
    et.run(G,1);
    rep(i,1,n)lis[a[i]].emb(i);
    rep(i,1,n)
    {
        if(lis[i].empty())continue;
        rep(j,1,T.etot)T.head[T.to[j]]=0;
        T.etot=0;
        auto r = ConciseTree::build(lis[i],T,db,et);
        dfs(T,r,0);
    }
    printf("%lld",ans*ans);
    return 0;
}