LeetCode 2421. 好路径的数目

LeetCode 2421. 好路径的数目

并查集 + 最小生成树思想

class Solution {
public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        int n = vals.size();
        vector<int> p(n);
        vector<int> cnt(n, 1);
        vector<int> q(n);
        vector<vector<int>> g(n);
        iota(p.begin(), p.end(), 0);
        iota(q.begin(), q.end(), 0);
        for(auto& e : edges)
        {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        function<int(int)> find;
        find = [&](int x) -> int {
            if(p[x] != x) p[x] = find(p[x]);
            return p[x];
        };

        sort(q.begin(), q.end(), [&](int a, int b){
            return vals[a] < vals[b];
        });

        int res = n;
        for(int i = 0; i < n; i ++)
        {
            int x = q[i];
            for(auto y : g[x])
                if(vals[x] >= vals[y])
                {
                    int a = find(x), b = find(y);
                    if(a != b && vals[a] == vals[b])
                    {
                        res += cnt[a] * cnt[b]; //树中的好路径数量
                        cnt[a] += cnt[b];
                    }
                    p[b] = a;
                }
        }
        return res;
    }
};