深搜+回溯，LeetCode 1766. 互质树

一、题目

1、题目描述

给你一个 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 。

2、接口描述

python3

​

class Solution:
    def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:

cpp

​

class Solution {
public:
    vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {

    }
};

3、原题链接

1766. 互质树

---

二、解题报告

1、思路分析

考虑数结点值范围只有50，我们可以预处理出所有值的质数

然后深搜，入，查，回，离

入：先找遍历过的互质数结点，记录深度最大那个

更新当前结点值的深度和结点信息，tmp保存旧信息，用于回溯恢复

查：查子节点

回：恢复现场

离：返回上一层

2、复杂度

时间复杂度： O(nU)空间复杂度：O(n  + U)，U = max(nums) = 50

3、代码详解

python3

​

N = 51
primes = [[j for j in range(1, N) if gcd(i, j) == 1] for i in range(N)]
class Solution:
    def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:
        n = len(nums)
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)
        res = [0] * n
        dep_id = [(-1, -1)] * N
        def dfs(x: int, fa: int, d: int)->None:
            v = nums[x]
            res[x] = max(dep_id[i] for i in primes[v])[1]
            tmp = dep_id[v]
            dep_id[v] = (d, x)
            for y in g[x]:
                if y != fa:
                    dfs(y, x, d + 1)
            dep_id[v] = tmp
        dfs(0, -1, 0)
        return res

cpp

​

typedef pair<int, int> PII;
const int N = 51;
vector<int> primes[N];
auto init = []{
    for(int i = 1; i < N; i++)
        for(int j = 1; j < N; j++)  
            if(__gcd(i, j) == 1)
                primes[i].emplace_back(j);
    return 0;
}();

class Solution {
public:
vector<int> res;
PII dep_id[N];
vector<vector<int>> g;

    void dfs(int x, int fa, int d, vector<int>& nums){
        int v = nums[x];
        int mad = 0;
        for(int y : primes[v]){
            auto [dep, id] = dep_id[y];
            if(dep > mad) mad = dep, res[x] = id;
        }
        auto tmp = dep_id[v];
        dep_id[v] = { d, x };
        for(int y : g[x])
            if(y != fa) dfs(y, x, d + 1, nums);
        dep_id[v] = tmp;
    }
    vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
        int n = nums.size();
        g.resize(n);res.resize(n, -1);
        for(auto& e : edges)
            g[e[0]].emplace_back(e[1]), g[e[1]].emplace_back(e[0]);
        dfs(0, -1, 1, nums);
        return res;
    }
};