2421. Number of Good Paths - Explanation

class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        n = len(vals)
        adj = [[] for _ in range(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)

        def dfs(node, startNode, parent):
            if vals[node] > vals[startNode]:
                return 0

            res = 0
            if vals[node] == vals[startNode] and node >= startNode:
                res += 1

            for child in adj[node]:
                if child == parent:
                    continue
                res += dfs(child, startNode, node)

            return res


        res = 0
        for node in range(n):
            res += dfs(node, node, -1)
        return res

class Solution:
    def numberOfGoodPaths(self, vals, edges):
        n = len(vals)
        adj = [[] for _ in range(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)

        res = 0
        for startNode in range(n):
            q = deque([startNode])
            visited = set([startNode])
            count = 0

            while q:
                node = q.popleft()
                if vals[node] == vals[startNode] and node >= startNode:
                    count += 1

                for child in adj[node]:
                    if child not in visited and vals[child] <= vals[startNode]:
                        visited.add(child)
                        q.append(child)

            res += count

        return res

class DSU:
    def __init__(self, n):
        self.Parent = list(range(n + 1))
        self.Size = [1] * (n + 1)

    def find(self, node):
        if self.Parent[node] != node:
            self.Parent[node] = self.find(self.Parent[node])
        return self.Parent[node]

    def union(self, u, v):
        pu, pv = self.find(u), self.find(v)
        if pu == pv:
            return False
        if self.Size[pu] >= self.Size[pv]:
            self.Size[pu] += self.Size[pv]
            self.Parent[pv] = pu
        else:
            self.Size[pv] += self.Size[pu]
            self.Parent[pu] = pv
        return True

class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        adj = collections.defaultdict(list)
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)

        valToIndex = collections.defaultdict(list)
        for i, val in enumerate(vals):
            valToIndex[val].append(i)

        res = 0
        uf = DSU(len(vals))

        for val in sorted(valToIndex.keys()):
            for i in valToIndex[val]:
                for nei in adj[i]:
                    if vals[nei] <= vals[i]:
                        uf.union(nei, i)

            count = collections.defaultdict(int)
            for i in valToIndex[val]:
                count[uf.find(i)] += 1
                res += count[uf.find(i)]

        return res

class DSU:
    def __init__(self, n, vals):
        self.parent = list(range(n))
        self.vals = vals
        self.count = [1] * n # count of nodes with max value of the component

    def find(self, node):
        if self.parent[node] != node:
            self.parent[node] = self.find(self.parent[node])
        return self.parent[node]

    def union(self, u, v):
        pu, pv = self.find(u), self.find(v)
        if pu == pv:
            return 0
        if self.vals[pu] < self.vals[pv]:
            self.parent[pu] = pv
        elif self.vals[pu] > self.vals[pv]:
            self.parent[pv] = pu
        else:
            self.parent[pv] = pu
            result = self.count[pu] * self.count[pv]
            self.count[pu] += self.count[pv]
            return result

        return 0


class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        n = len(vals)
        dsu = DSU(n, vals)

        # Sort edges based on max value of the two nodes
        edges.sort(key=lambda edge: max(vals[edge[0]], vals[edge[1]]))

        res = n  # Each node alone is a good path
        for u, v in edges:
            res += dsu.union(u, v)
        return res