一道虛樹題

題目大意

給樹上每個點規定一種顏色

問同色距離最大的點對距離平方是多少

虛樹

每次給定樹上一個點集(大小爲$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;
}