LC 1766. 互质树

题目描述

这是 LeetCode 上的 1766. 互质树 ，难度为 困难。

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

给你一个整数数组 nums 和一个二维数组 edges 来表示这棵树。

nums[i] 表示第 i 个点的值，$edges[j] = [u{j}, v{j}]$ 表示节点 $u{j}$ 和节点 $v{j}$ 在树中有一条边。

当 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：

输入：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]

提示：

• $nums.length = n$
• $1 <= nums[i] <= 50$
• $1 <= n <= 10^5$
• $edges.length = n - 1$
• $edges[j].length = 2$
• $0 <= u{j}, v{j} < n$
• $u{j} \neq v{j}$

---

DFS

题目描述很长，但其实就是说每个节点从下往上找，找到最近的「与其互质」的节点。

数据范围是 $10^5$，如果每个节点都直接往上找最近「互质」祖宗节点的话，当树为线性时，复杂度是 $O(n^2)$ ，会超时。

因此我们要利用 $nums[i]$ 范围只有 $50$ 的特性。

我们可以先预处理除 $[1, 50]$ 范围内的每个数，求出他们互质的数有哪些，存到一个字典里。

那么对于某个节点而言，假设节点的值为 x ，所在层数为 y。

那么问题转化为求与 x 互质的数有哪些，最近的在哪一层。

用 dep[x] 表示距离值为 x 的节点最近的层是多少；pos[x] 代表具体的节点编号。

代码：

class Solution {
    int[] ans;
    Map<Integer, List<Integer>> map = new HashMap<>(); // 边映射
    Map<Integer, List<Integer>> val = new HashMap<>(); // 互质数字典
    int[] dep;
    int[] pos = new int[52];
    public int[] getCoprimes(int[] nums, int[][] edges) {
        int n = nums.length;
        ans = new int[n];
        dep = new int[n];
        Arrays.fill(ans, - 1);
        Arrays.fill(pos, -1);

        for (int[] edge : edges) {
            int a = edge[0], b = edge[1];
            List<Integer> alist = map.getOrDefault(a, new ArrayList<>());
            alist.add(b);
            map.put(a, alist);
            List<Integer> blist = map.getOrDefault(b, new ArrayList<>());
            blist.add(a);
            map.put(b, blist);
        }

        for (int i = 1; i <= 50; i++) {
            for (int j = 1; j <= 50; j++) {
                if (gcd(i, j) == 1) {
                    List<Integer> list = val.getOrDefault(i, new ArrayList<>());
                    list.add(j);
                    val.put(i, list);
                }
            }
        }

        dfs(nums, 0, -1);
        return ans;
    }
    void dfs(int[] nums, int u, int form) {
        int t = nums[u];
        for (int v : val.get(t)) {
            if (pos[v] == -1) continue;
            if (ans[u] == -1 || dep[ans[u]] < dep[pos[v]]) ans[u] = pos[v];
        }
        int p = pos[t];
        pos[t] = u;

        for (int i : map.get(u)) {
            if (i == form) continue;
            dep[i] = dep[u] + 1;
            dfs(nums, i, u);
        }
        pos[t] = p;
    }
    int gcd(int a, int b) {
        if (b == 0) return a;
        if (a == 0) return b;
        return gcd(b, a % b);
    }
}

• 时间复杂度：对于每个节点而言，会检查与其数值互质的数有哪些，在哪层。最坏情况下会检查 50 个互质数（当前数值为 1）。复杂度为 $O(n)$
• 空间复杂度：$O(n)$

---

最后

这是我们「刷穿 LeetCode」系列文章的第 No.1766 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。