题解 | #树上最短链# 邻接表和bfs

#include <iostream>
#include <algorithm>
#include <cmath>
#include <queue>
#include <cstring>
using namespace std;
//无向边存两遍
const int N = 10010;
int n;
//城市的等级
int A[N];

//邻接表
int e[N];
int ne[N];
int h[N];
int idx;

typedef pair<int,int> ii;

void add(int a,int b)
{
    e[idx] = b;
    ne[idx] = h[a];
    h[a] = idx++;
}

int minL(int k){
    
    int st[N]={0};
    queue<ii> q;
    //当前节点入队
    q.push({k,0});
    st[k] = 1;
    //队列不空，队头出队
    while(q.size())
    {
        ii t = q.front();
        q.pop();
        //cout << "出队：" << t.first << " --- "<< t.second << endl;
        if(t.first!=k && A[t.first]==A[k])
            return t.second;
        //扩展
        int i = h[t.first];
        while(i!=-1)
        {
            int j = e[i];
            if(st[j]==0)
            {
                q.push({j,t.second+1});
                st[j] = 1;
            }
            i = ne[i];
        }
    }
    return -1;
}

int main()
{
    memset(h,-1,sizeof h);
    scanf("%d", &n);
    for (int i = 1;i <= n; i++)
    {
        scanf("%d",&A[i]);
    }
    for(int i = 1;i<=n-1;i++)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        add(a,b);
        add(b,a);
    }
    //给每个节点 找最近的同级别的节点，保存路径长度
    int ans = 1e9;
    for(int i = 1;i<=n;i++)
    {
        int tmp = minL(i);
        if(tmp==-1)
            continue;
        if(tmp <ans)
            ans = tmp;
    }
    if(ans==1e9)
        printf("%d",-1);
    else
        printf("%d",ans);
    return 0;
}