Number of Good Paths

class DisjointSet {
private:
    vector<int> leader, size;
    
public:
    DisjointSet(int cnt) {
        leader.resize(cnt);
        size.resize(cnt, 0);
        for (int id = 0; id < cnt; id++) {
            leader[id] = id;
        }
    }
    int find_leader(int x) {
        if (leader[x] != x) leader[x] = find_leader(leader[x]);
        return leader[x];
    }
    void merge(int x, int y) {
        int rootX = find_leader(x), rootY = find_leader(y);
        if (rootX == rootY) {
            return;
        } else if (size[rootX] < size[rootY]) {
            leader[rootX] = rootY;
        } else if (size[rootX] > size[rootY]) {
            leader[rootY] = rootX;
        } else {
            leader[rootY] = rootX;
            size[rootX]++;
        }
    }
};
    
class Solution {
public:
    int numberOfGoodPaths(vector<int>& values, vector<vector<int>>& links) {
        int nodeCount = values.size();
        vector<vector<int>> adjacency(nodeCount);
        for (auto& link : links) {
            adjacency[link[0]].push_back(link[1]);
            adjacency[link[1]].push_back(link[0]);
        }
        // Group nodes by values in ascending order.
        map<int, vector<int>> valueGroupMap;
        for (int node = 0; node < nodeCount; node++) {
            valueGroupMap[values[node]].push_back(node);
        }
    
        DisjointSet dsu(nodeCount);
        int pathCount = 0;
    
        // Process each group of nodes with the same value.
        for (auto& [val, nodes] : valueGroupMap) {
            // Link sets of the node with its neighbors that have equal or less value.
            for (int node : nodes) {
                for (int neighbor : adjacency[node]) {
                    if (values[node] >= values[neighbor]) {
                        dsu.merge(node, neighbor);
                    }
                }
            }
            // Count all nodes in the current sets to calculate number of paths.
            unordered_map<int, int> componentSize;
            for (int u : nodes) {
                int root = dsu.find_leader(u);
                componentSize[root]++;
            }
            // Calculate paths for each group of connected nodes.
            for (auto& [_, size] : componentSize) {
                pathCount += (size * (size + 1) / 2);
            }
        }
        return pathCount;
    }
};

class DisjointSet {
    int[] leader;
    int[] treeRank;
    
    public DisjointSet(int capacity) {
        leader = new int[capacity];
        for (int i = 0; i < capacity; i++)
            leader[i] = i;
        treeRank = new int[capacity];
    }
    
    public int findRoot(int x) {
        if (leader[x] != x)
            leader[x] = findRoot(leader[x]);
        return leader[x];
    }
    
    public void mergeSets(int x, int y) {
        int rootX = findRoot(x), rootY = findRoot(y);
        if (rootX == rootY) {
            return;
        } 
        if (treeRank[rootX] < treeRank[rootY]) {
            leader[rootX] = rootY;
        } else if (treeRank[rootX] > treeRank[rootY]) {
            leader[rootY] = rootX;
        } else {
            leader[rootY] = rootX;
            treeRank[rootX]++;
        }
    }
}
    
class Solution {
    public int numberOfGoodPaths(int[] vals, int[][] edges) {
        Map<Integer, List<Integer>> graph = new HashMap<>();
        for (int[] edge : edges) {
            int u = edge[0], v = edge[1];
            graph.computeIfAbsent(u, k -> new ArrayList<Integer>()).add(v);
            graph.computeIfAbsent(v, k -> new ArrayList<Integer>()).add(u);
        }
    
        int n = vals.length;
        TreeMap<Integer, List<Integer>> valueGroups = new TreeMap<>();
        for (int i = 0; i < n; i++) {
            valueGroups.computeIfAbsent(vals[i], k -> new ArrayList<Integer>()).add(i);
        }
    
        DisjointSet dsu = new DisjointSet(n);
        int totalPaths = 0;
    
        for (int value : valueGroups.keySet()) {
            for (int node : valueGroups.get(value)) {
                if (!graph.containsKey(node))
                    continue;
                for (int neighbor : graph.get(node)) {
                    if (vals[node] >= vals[neighbor]) {
                        dsu.mergeSets(node, neighbor);
                    }
                }
            }
            Map<Integer, Integer> countInSet = new HashMap<>();
            for (int u : valueGroups.get(value)) {
                countInSet.put(dsu.findRoot(u), countInSet.getOrDefault(dsu.findRoot(u), 0) + 1);
            }
            for (int setSize : countInSet.values()) {
                totalPaths += setSize * (setSize + 1) / 2;
            }
        }
        return totalPaths;
    }
}