算法leetcode

力扣1766-互质树

字数统计: 1.4k字   |   阅读时长: 5分
, , ,

摘要: 祖先链的信息，树形前缀

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

各位好，今天我们继续来看一个祖先链和树形前缀的题，在 DFS 的过程中记录祖先链上的质数因子相关的信息。之前的一些文章中我们讨论过祖先链和树形前缀的问题，思路类似，可以参考：

题目

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

给你一个整数数组 nums 和一个二维数组 edges 来表示这棵树。nums[i] 表示第 i 个点的值，edges[j] = [uj, vj] 表示节点 uj 和节点 vj 在树中有一条边。

当 gcd(x, y) == 1 ，我们称两个数 x 和 y 是 互质的 ，其中 gcd(x, y) 是 x 和 y 的 最大公约数 。

从节点 i 到 根 最短路径上的点都是节点 i 的祖先节点。一个节点 不是 它自己的祖先节点。

请你返回一个大小为 n 的数组 ans ，其中 ans[i]是离节点 i 最近的祖先节点且满足 nums[i] 和 nums[ans[i]] 是 互质的 ，如果不存在这样的祖先节点，ans[i] 为 -1 。

提示：

1
2
3
4
5
6
7
nums.length == n
1 <= nums[i] <= 50
1 <= n <= 1e5
edges.length == n - 1
edges[j].length == 2
0 <= uj, vj < n
uj != vj

示例 1：

输入：nums = [2,3,3,2], edges = [[0,1],[1,2],[1,3]]
输出：[-1,0,0,1]
解释：上图中，每个节点的值在括号中表示。

  • 节点 0 没有互质祖先。
  • 节点 1 只有一个祖先节点 0 。它们的值是互质的（gcd(2,3) == 1）。
  • 节点 2 有两个祖先节点，分别是节点 1 和节点 0 。节点 1 的值与它的值不是互质的（gcd(3,3) == 3）但节点 0 的值是互质的(gcd(2,3) == 1)，所以节点 0 是最近的符合要求的祖先节点。
  • 节点 3 有两个祖先节点，分别是节点 1 和节点 0 。它与节点 1 互质（gcd(3,2) == 1），所以节点 1 是离它最近的符合要求的祖先节点。

示例 2：

输入：nums = [5,6,10,2,3,6,15], edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]
输出：[-1,0,-1,0,0,0,-1]

题解

算法: DFS、祖先链、树形前缀

对于节点 $u$，我们要找的是 $u$ 的祖先链上，顶点权重与 $u$ 的权重互质的顶点中最深（也就是离 $u$ 最近）的那一个。

算法的大框架是对树做一次 DFS，访问到节点 $u$ 时，$u$ 的祖先链上的节点已经访问过，此时做下面两件事：

  • 根据祖先链中的信息，更新当前节点的答案 $result[u]$。
  • 将当前节点 $u$ 的信息记录。

DFS 算法可以高效运行的核心是如何更新信息与答案。最暴力的做法是对于 $u$，从其父节点 $fa$ 开始，从深到浅依次访问祖先链上的各个节点 $v$，看 $nums[u]$ 与 $nums[v]$ 是否互质。但访问祖先链最坏也要 $O(N)$，算法的总时间复杂度就来到 $O(N^{2})$。

因此我们必须记录某种树形前缀信息 $pre$，使得我们在访问到 $u$ 时，不用依次访问祖先链各个节点，只需要从树形前缀信息 $pre$ 中的记录，即可高效计算出答案 $result[u]$。

由于顶点权重的取值范围是 $1 \sim 50$，这个范围非常重要，使得我们可以对每一个可能值维护一个前缀信息。

记 $pre[p][u]$ 表示如果顶点 $u$ 的权值为 $p$，则在 $u$ 的祖先链中，权值与 $nums[u]$ 互质且距离最近的顶点编号。顶点权重范围为 $1 \sim 50$，这也是 $p$ 的范围。

如果有这样的信息，DFS 访问到 $u$ 点时计算答案 $result[u]$ 就非常方便：

下一个问题就是如何高效更新 $pre$。访问到 $u$ 时，依次枚举 $p = 1, 2, \cdots, 50$，做以下更新：

注意初始化的问题，也就是根节点的 $pre$ 信息为 $pre[p][0] = -1$。

代码 (Python)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Solution:
    def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:
        n = len(nums)
        g = [[] for _ in range(n)]
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)

        # pre[p][u] := 从根到 fa 的这段前缀上，距离 u 最近且与 p 互质的节点
        pre = [[-1 for _ in range(n)] for _ in range(51)]
        result = [-1] * n

        def dfs(u: int, fa: int):
            if fa != -1:
                # 计算 result[u]
                if math.gcd(nums[u], nums[fa]) == 1:
                    result[u] = fa
                else:
                    result[u] = pre[nums[u]][fa]
                # 推导 pre[p][u]
                for p in range(1, 51):
                    if math.gcd(p, nums[fa]) == 1:
                        pre[p][u] = fa
                    else:
                        pre[p][u] = pre[p][fa]

            for v in g[u]:
                if v == fa:
                    continue
                dfs(v, u)

        dfs(0, -1)
        return result
Share