树上BFS问题

题意简述：

n个城市有n-1条无向边的连通图（也是树结构），每条边的时间权值都是1； 每个城市有一个等级，求以等级相同的任意两城市作为起点和终点时，最小的时间花费。（每条边至多经过一次）

算法设计

我们可以以任意结点为起点，以其他任意结点为终点，基本上就是《任意两点间的最短路径问题》，只是多了一个限制条件，起点和终点的等级相同，但本质上还是要遍历每一个结点才能成功到达任意一个结点。所以算法复杂度的期望也就只能达到O(N * V)，即节点数乘以边数，也就是O(n^2)

那我们枚举每一个起点，找出离它最近的第一个相同等级的终点，明显使用 bfs 较快。 因为是树结构，我们可以直接以起点为根节点展开成树，到达每个结点的时间花费就是其深度，那么每次 bfs 直接返回深度最小的那个相同等级结点的深度即可，若没有同等级结点则返回正无穷。

一次 bfs 时间复杂度为 O(n)，对每个结点都 bfs 一次，则 n 次 bfs 的时间复杂度为 O(n^2)，任务完成。

C++代码

#include<vector>
#include<algorithm>
#include<queue>
#include<set>
using namespace std;

// 计算以st为起点的最短路径，以st结点为根节点展开成树
    int No6_bfs(int st, const vector<vector<int>>& to, const vector<int>& level) {
        queue<int>Q;
        Q.push(st);
        vector<bool>vis(to.size(), false);
        vis[st] = true;
        int cur = 0;
        while (!Q.empty()) {
            cur++;  //连通图不存在孤立点，一定可以走一步
            //下面把这一层全部搜索一遍
            int k = Q.size();
            while (k--) {
                int t = Q.front();
                Q.pop();
                for (int x : to[t]) {   // t -> x 的边
                    if (vis[x])continue;
                    if (level[x] == level[st]) { //起点与终点的等级相同
                        return cur;
                    }
                    Q.push(x);
                    vis[x] = true;
                }
            }
        }
        return 1 << 30;
    }
    int No6_Tree(vector<int>& level, vector<pair<int, int>>& edge) {
        int n = level.size();   // 注意edge中的结点编号是从1到n
        vector<vector<int>> to(n);  //邻接表中的结点编号从0到n-1
        for (int i = 0; i < n - 1; i++) {
            to[edge[i].first - 1].push_back(edge[i].second - 1);
            to[edge[i].second - 1].push_back(edge[i].first - 1);
        }
        int ans(1 << 30);
        for (int i = 0; i < n; i++) {   //枚举每个起点
            // 下面 bfs 搜索出 level[i]等级的最短路径
            ans = min(ans, No6_bfs(i, to, level));
        }
        //n个城市等级从小到大
        return ans == (1 << 30) ? -1 : ans;
    }

int main(){
    int n;
    cin>>n;
    vector<int> level(n);
    vector<pair<int, int>> edge(n-1);
    for(int i=0;i<n;i++)
        cin>>level[i];
    for(int i=0;i<n-1;i++){
        cin>>edge[i].first>>edge[i].second;
    }
    cout<<No6_Tree(level, edge)<<'\n';
}