HihoCoder - 1387 求树上相同颜色的直径

题目： 给你一棵n个节点的树，每个节点的颜色可能不同，现在要给你两种颜色，问你两种颜色的最大距离。如果有一种颜色不存在那么直接输出-1 即可。

#include<bits/stdc++.h>
using namespace std;
const int inf=99999999;
const int maxn=1e5+10;
int n,m,a[maxn];
struct node{
    int v,nxt;
}edge[maxn<<1];
int first[maxn],tot;
void addedge(int u,int v)
{
    tot++;
    edge[tot].v=v;
    edge[tot].nxt=first[u];
    first[u]=tot;
}
int dep[maxn],f[maxn][22];
void dfs(int u,int fa)
{
    for(int i=first[u];i!=-1;i=edge[i].nxt)
    {
        int v=edge[i].v;
        if(v==fa) continue;
        dep[v]=dep[u]+1;
        f[v][0]=u;
        dfs(v,u);
    }
}

int dist(int x,int y)
{
    int ans=dep[x]+dep[y];
    if(dep[x]<dep[y])
        swap(x,y);
    for(int i=20;i>=0;i--)
    {
        int k=f[x][i];
        if(dep[k]>=dep[y])
            x=k;
    }
    if(x!=y)
    {
        for(int i=20;i>=0;i--)
        {
            if(f[x][i]!=f[y][i])
            {
                x=f[x][i];
                y=f[y][i];
            }
        }
        x=f[x][0];
    }
    ans-=2*dep[x];
    return ans+1;
}
map<string,int>ma;
string s,s1,s2;
char ch[1000];
int seg[maxn][2];
int main()
{
    while(~scanf("%d%d",&n,&m))
    {
        int id=0;
        ma.clear();
        for(int i=1;i<=n;i++)
        {
            scanf("%s",ch);
            s=(string)ch;
            if(!ma[s])
                ma[s]=++id;
            a[i]=ma[s];
        }
        tot=0;
        for(int i=1;i<=n;i++)
            first[i]=-1;
        for(int i=1;i<n;i++)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            addedge(x,y);
            addedge(y,x);
        }
        dep[1]=1;
        dfs(1,-1);
        for(int j=1;j<=20;j++)
            for(int i=1;i<=n;i++)
                f[i][j]=f[f[i][j-1]][j-1];
        memset(seg,0,sizeof seg);

        for(int i=1;i<=n;i++)
        {
            int k=a[i];
            if(!seg[k][0])
            {
                seg[k][0]=seg[k][1]=i;
                continue;
            }
            int u=seg[k][0],v=seg[k][1];
            int len=dist(u,v);
            int tmp1=dist(u,i),tmp2=dist(v,i);
            if(len<tmp1)
            {
                seg[k][0]=u;
                seg[k][1]=i;
                len=tmp1;
            }
            if(len<tmp2)
            {
                seg[k][0]=v;
                seg[k][1]=i;
            }
        }
        while(m--)
        {
            scanf("%s",ch);
            s1=(string)ch;
            scanf("%s",ch);
            s2=(string)ch;
            int x=ma[s1],y=ma[s2];
            if(x==0||y==0)
            {
                puts("-1");
                continue;
            }
            int u[3],v[3];
            u[0]=seg[x][0];u[1]=seg[x][1];
            v[0]=seg[y][0];v[1]=seg[y][1];
            int ans=0;
            for(int i=0;i<2;i++)
                for(int j=0;j<2;j++)
                ans=max(ans,dist(u[i],v[j]));
            printf("%d\n",ans);
        }
    }
    return 0;
}

一开始用cin读入时，加上ios::sync_with_stdio(0);后就WA了。。。