好路径的数目（并查集，启发式）

题意

给定一颗树，有点权，问好路径的数目（unordered_pair）。
好路径定义：

• 开始节点和结束节点点权相同
• 开始节点权值为路径点权max

constrains:
$1\le n\le 3e4$
$1\le va{l}_i\le 1e5$

思路

考虑按点权从小到大重新构建这棵树（可能见过这种思维，但比赛中完全没考虑），已插入的点之间连边，需要维护并查集，以及该连通块中含有当前最大权值的点的个数，从而更新答案。
启发式合并也能做。

代码

class Solution {
public:
    map<int, vector<int>> G;
    vector<int> va;
    vector<int> ad[30010];
    int f[30010];
    map<int,int> cnt;
    bool c[30010];   
    int n,ans;
    int findfa(int x){
        return x==f[x]?x:f[x]=findfa(f[x]);
    }
    void merge(int u,int v){
        u = findfa(u);v = findfa(v);
        if(u!=v){
            f[v] = u;
            ans+=cnt[u]*cnt[v];
            cnt[u] += cnt[v];
        }
    }
    
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        va = vals;
        ans = 0;
        n = vals.size();
        for(int i = 0;i<vals.size();i++){
            G[vals[i]].push_back(i);
        }
        for(int i = 0;i<n;i++) f[i] = i;
        for(auto s:edges){
            ad[s[0]].push_back(s[1]);
            ad[s[1]].push_back(s[0]);
        }
        for(auto [a,b]:G){
            cnt.clear();
            for(auto u:b){
                c[u] = 1;
                cnt[u]++;
                for(auto v:ad[u]){
                    if(!c[v]) continue;
                    merge(u,v);
                }
            }
        }
        return ans+n;
    }
};

启发式合并

class Solution {
    typedef pair<int, int> pii;

    vector<int> A;
    long long ans = 0;
    vector<vector<int>> e;

    map<int, int> dfs(int sn, int fa) {
        // 返回的 map 里一开始只有子树的根
        map<int, int> ret;
        ret[A[sn]]++;

        for (int fn : e[sn]) if (fn != fa) {
            // 递归调用
            map<int, int> tmp = dfs(fn, sn);
            // 清除子节点里权值不符合要求的路径。
            // 这里好像枚举了 tmp，破坏了启发式合并的复杂度，
            // 但由于我们枚举到的每个元素都会被删掉，所以总共最多只有 n 次删除，复杂度还是正确的。
            for (auto it = tmp.begin(); it != tmp.end() && it->first < A[sn]; ) tmp.erase(it++);
            // 启发式合并的关键，每次只能枚举较小的结果（复杂度的原理等同于轻重链剖分https://oi-wiki.org/graph/dsu-on-tree/）
            if (tmp.size() > ret.size()) swap(tmp, ret);
            // 统计路径中深度最小的点为 sn 的好路径数量
            for (auto it = tmp.begin(); it != tmp.end(); it++)
                if (ret.count(it->first)) ans += (it->second) * ret[it->first];
            // 更新返回值
            for (auto it = tmp.begin(); it != tmp.end(); it++) ret[it->first] += it->second;
        }

        return ret;
    }

public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        int n = vals.size();
        A = vals;
        e.resize(n);
        for (auto &vec : edges) {
            e[vec[0]].push_back(vec[1]);
            e[vec[1]].push_back(vec[0]);
        }

        dfs(0, 0);
        return ans + n;
    }
};

作者：TsReaper
链接：https://leetcode.cn/circle/discuss/f2jB1Y/view/b4qNEp/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。