好路径的数目

标签: 树 并查集 图 数组 哈希表 排序

难度: Hard

给你一棵 n 个节点的树（连通无向无环的图），节点编号从 0 到 n - 1 且恰好有 n - 1 条边。

给你一个长度为 n 下标从 0 开始的整数数组 vals ，分别表示每个节点的值。同时给你一个二维整数数组 edges ，其中 edges[i] = [ai, bi] 表示节点 ai 和 bi 之间有一条 无向 边。

一条 好路径 需要满足以下条件：

  1. 开始节点和结束节点的值 相同 。
  2. 开始节点和结束节点中间的所有节点值都 小于等于 开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。

请你返回不同好路径的数目。

注意，一条路径和它反向的路径算作 同一 路径。比方说， 0 -> 1 与 1 -> 0 视为同一条路径。单个节点也视为一条合法路径。

示例 1：

输入：vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]
输出：6
解释：总共有 5 条单个节点的好路径。
还有 1 条好路径：1 -> 0 -> 2 -> 4 。
（反方向的路径 4 -> 2 -> 0 -> 1 视为跟 1 -> 0 -> 2 -> 4 一样的路径）
注意 0 -> 2 -> 3 不是一条好路径，因为 vals[2] > vals[0] 。

示例 2：

输入：vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]
输出：7
解释：总共有 5 条单个节点的好路径。
还有 2 条好路径：0 -> 1 和 2 -> 3 。

示例 3：

输入：vals = [1], edges = []
输出：1
解释：这棵树只有一个节点，所以只有一条好路径。

提示：

  • n == vals.length
  • 1 <= n <= 3 * 104
  • 0 <= vals[i] <= 105
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • ai != bi
  • edges 表示一棵合法的树。

Submission

运行时间: 246 ms

内存: 29.4 MB

class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        n = len(vals)
        parent = [i for i in range(n)]
        maxv = [vals[i] for i in range(n)]
        cnt = [1] * n
        res = n

        def find(x):
            p = x
            while parent[p] != p:
                p = parent[p]
            while parent[x] != p:
                parent[x], x = p, parent[x] 
            return p
        
        def union(x, y):
            nonlocal res
            p1 = find(x)
            p2 = find(y)
            if p1 != p2:
                if maxv[p1] > maxv[p2]:
                    parent[p2] = p1
                elif maxv[p1] < maxv[p2]:
                    parent[p1] = p2
                else:
                    res += cnt[p1] * cnt[p2]
                    cnt[p1] += cnt[p2]
                    parent[p2] = p1

        edges.sort(key=lambda x: max(vals[x[0]], vals[x[1]]))
        for x, y in edges:
            union(x, y)
        return res

Explain

这个题解使用了一个联合查找(Union-Find)数据结构，用来高效地处理节点的连接和根查找。基本思路是：1. 初始化每个节点为独立的集合。2. 对边进行排序，以边两端节点的最大值为键进行排序。这样可以保证当处理一条边时，路径上较大的值已经处理过，符合好路径定义中的“中间所有节点值都小于等于起始节点值”。3. 通过并查集合并操作来连接节点，如果两个节点属于不同集合且值相同，则这两个集合可以合并，并且此时这两个集合可能的路径数为两个集合大小的乘积。4. 遍历所有边，进行合并操作，并计算可能的好路径。单个节点的好路径（节点自身）在初始化时就加入了总数。

时间复杂度: O(n log n)

空间复杂度: O(n)

# 加注释的题解代码

class Solution:
    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:
        n = len(vals)
        parent = [i for i in range(n)] # 初始化并查集的父节点指向自己
        maxv = [vals[i] for i in range(n)] # 记录与并查集代表元素相同集合的最大值
        cnt = [1] * n # 记录集合的大小
        res = n # 所有单节点的路径数

        def find(x):
            p = x
            while parent[p] != p:
                p = parent[p]
            while parent[x] != p:
                parent[x], x = p, parent[x]  # 路径压缩
            return p
        
        def union(x, y):
            nonlocal res
            p1 = find(x)
            p2 = find(y)
            if p1 != p2:
                if maxv[p1] > maxv[p2]:
                    parent[p2] = p1
                elif maxv[p1] < maxv[p2]:
                    parent[p1] = p2
                else:
                    res += cnt[p1] * cnt[p2] # 两个集合可以合并，增加路径数
                    cnt[p1] += cnt[p2]
                    parent[p2] = p1

        edges.sort(key=lambda x: max(vals[x[0]], vals[x[1]])) # 按边的最大值排序
        for x, y in edges:
            union(x, y)
        return res

Explore

按最大节点值对边进行排序的目的是为了保证处理每条边时，较大值的节点已经被处理过了。这样的排序策略可以确保在并查集中合并节点时，已经考虑了所有可能连接到当前边的较小值节点，从而满足好路径定义即路径上的所有中间节点值都小于等于起始节点值的要求。这种排序帮助算法按照从小到大的顺序处理节点值相同的一组节点，简化了路径计数的逻辑，因为在合并时可以确保不会引入比当前节点值大的节点，避免了额外的复杂度。

`maxv`数组在并查集中记录了每个集合的最大节点值。它的维护发生在`union`函数中，当两个集合合并时，需要更新新的集合的最大值。`maxv`数组确保在合并过程中，只有当两个集合的最大值相同时，才会进行合并操作，这样可以保证合并的集合中所有节点的值不会超过这个最大值。这对于满足好路径的条件至关重要，即所有路径上的节点值需满足不大于路径起点和终点的值。

在`union`函数中比较`maxv[p1]`和`maxv[p2]`主要是为了确保合并的两个集合符合好路径的条件，即集合内的任何节点值都不超过集合中的最大值。如果两个集合的最大值不同，则不应该合并，因为这可能导致合并后的集合包含不符合好路径条件的节点。这种比较确保了只有在两个集合的最大值相等时才进行合并，从而控制了合并后集合的属性，保证了算法的正确性和高效性。

在题解中`res += cnt[p1] * cnt[p2]`这一步的计算逻辑是在两个集合合并时，计算由这两个集合的所有可能组合形成的新路径数。因为在排序和合并过程中已经确保了路径值的单调性和集合的独立性，所以每次合并时计算的是新生成的路径，不会与之前的路径重复。此外，由于路径是按节点值从小到大处理并计算的，因此不会出现重复计算同一条路径的反向路径的情况，每条路径只会在两个集合首次合并时被计算，保证了计算的唯一性和正确性。