参考题解。
思路:根据权值从小到大利用并查集建树算方案。
class Solution {
public:
int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
int n=vals.size();
vector<vector<int>> g(n);
for( vector<int> v:edges ){
g[v[0]].push_back(v[1]);
g[v[1]].push_back(v[0]);
}
int id[n],fa[n],siz[n];
for( int i=0;i<n;i++ ) id[i]=fa[i]=i,siz[i]=1;
function<int(int)> find = [&](int x)->int{ return fa[x]==x ? x : fa[x]=find(fa[x]); };
int ans=n;
sort(id,id+n,[&](int x,int y){ return vals[x]<vals[y]; });
for(int x:id ){
int v = vals[x], fx=find(x);
for( int y:g[x] ){
y=find(y);
if( y==fx || vals[y]>v ) continue;
if( vals[y]==v ){
ans+=siz[y]*siz[fx];
siz[fx]+=siz[y];
}
fa[y]=fx;
}
}
return ans;
}
};