牛客—— 漂亮的公园 (离散化+LCA+结论)

原题链接

题意:

给定一棵树,每个节点都对应有一个颜色,问询问的两个颜色之间的最大距离(可两个颜色可以相同也可以不同)。

思路:

首先要知道一个结论是对于同一种颜色的直径端点a,b,另一个点c到达这种颜色的最大值,一定是dis(a,c)和dis(b,c)中的最大值maxx。所以对于两种颜色,只需要枚举两种颜色的直径端点,取最大值即可。

证明的话是看了tzc_zwj的博客,假设存在d点使得dis(d,c)大于maxx,当d在ab的路径上时,显然maxx更大,因为加上了路径上的一段;当d不在ab的路径上时,显然可以通过d到路径上,再到端点,也是maxx更大。

这样的话两种颜色的最大距离就转化成了枚举两种颜色的直径的端点求距离取max。

然后对于树上求两点之间距离,我们可以借助LCA求解。

dep[x]+dep[y]-2*dep[lca(x,y)];

但是因为颜色很大,直接求可能会爆时间和空间,所以我们要先将颜色离散化。

综上,我们要先将颜色离散化,存储一下每种颜色对应的节点,利用LCA预处理出每种颜色的直径端点,在询问的时候枚举计算取最值即可。

代码:

#pragma GCC optimize(2)
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int>PII;
#define I_int ll
inline ll read()
{
    ll x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
char F[200];
inline void out(I_int x) {
    if (x == 0) return (void) (putchar('0'));
    I_int tmp = x > 0 ? x : -x;
    if (x < 0) putchar('-');
    int cnt = 0;
    while (tmp > 0) {
        F[cnt++] = tmp % 10 + '0';
        tmp /= 10;
    }
    while (cnt > 0) putchar(F[--cnt]);
    cout<<endl;
}

const int maxn=1e5+100;

int n,q,a[maxn],c[maxn];
vector<int>col[maxn];
int h[maxn],e[maxn*2],ne[maxn*2],idx;
int dep[maxn],fa[maxn][20];
int d[maxn][2];
void add(int u,int v){
    e[idx]=v,ne[idx]=h[u],h[u]=idx++;
}
void bfs(){
    queue<int>q;
    memset(dep,0x3f,sizeof dep);
    dep[0]=0;///特殊处理
    int root=1;
    dep[root]=1;q.push(root);
    while(!q.empty()){
        int t=q.front();q.pop();
        for(int i=h[t];~i;i=ne[i]){
            int j=e[i];
            if(dep[j]>dep[t]+1){
                dep[j]=dep[t]+1;
                q.push(j);
                fa[j][0]=t;///预处理fa数组
                for(int k=1;k<=15;k++)
                    fa[j][k]=fa[fa[j][k-1]][k-1];
            }
        }
    }
}

int lca(int a,int b){
    if(dep[a]<dep[b]) swap(a,b);
    for(int k=15;k>=0;k--)
        if(dep[fa[a][k]]>=dep[b]) a=fa[a][k];///使a,b跳到同一层
    if(a==b) return a;
    for(int k=15;k>=0;k--)
        if(fa[a][k]!=fa[b][k]) a=fa[a][k],b=fa[b][k];
    return fa[a][0];
}
///树上求两点距离:LCA
int cul(int x,int y){
    return dep[x]+dep[y]-2*dep[lca(x,y)];
}

int main(){

    n=read();q=read();
    for(int i=1;i<=n;i++){
        a[i]=read();
        c[i]=a[i];
    }
    ///离散化
    sort(c+1,c+1+n);
    int m=unique(c+1,c+1+n)-(c+1);
    for(int i=1;i<=n;i++){
        a[i]=lower_bound(c+1,c+1+m,a[i])-(c+1)+1;
        col[a[i]].push_back(i);
    }
    ///存储边
    memset(h,-1,sizeof h);

    for(int i=1;i<n;i++){
        int u=read(),v=read();
        add(u,v);add(v,u);
    }
    ///预处理LCA
    bfs();
    ///预处理每种颜色距离最远的两个点
    for(int i=1;i<=m;i++){///枚举颜色
        d[i][0]=d[i][1]=col[i][0];
        int d1=cul(d[i][0],d[i][1]);
        for(int j=1;j<col[i].size();j++){
            int x=col[i][j];
            int d2=cul(d[i][0],x),d3=cul(d[i][1],x);
            if(d2<d3) swap(d2,d3),swap(d[i][0],d[i][1]);
            if(d2>d1) d[i][1]=col[i][j],d1=cul(d[i][0],d[i][1]);
        }
    }

    while(q--){
        int x=read(),y=read();
        int cx=lower_bound(c+1,c+1+m,x)-c,cy=lower_bound(c+1,c+1+m,y)-c;
        if(cx>m||cy>m||c[cx]!=x||c[cy]!=y){
            ///颜色不存在或不符合要求
            puts("0");
            continue;
        }
        int res=0;
        ///枚举求解
        for(int i=0;i<2;i++)
            for(int j=0;j<2;j++)
                res=max(res,cul(d[cx][i],d[cy][j]));
        cout<<res<<endl;
    }

    return 0;
}