LC每日一题｜20240411 - 1766. 互质树

LC每日一题｜20240411 - 1766. 互质树

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

提示：

• nums.length == n
• 1 <= nums[i] <= 50
• 1 <= n <= 10^5
• edges.length == n - 1
• edges[j].length == 2
• 0 <= uj, vj < n
• uj != vj

题目等级：Hard

前置知识

求取最大公约数，我们需要用到 欧几里得算法，参见 OI-WIKI - 最大公约数 。

它大概长这样，建议背诵全文：

fun gcd(a: Int, b: Int): Int {
    if (b == 0) return a
    return gcd(b, a % b)
}

解题思路

我们发现 nums[i] 的取值范围非常小，所以我们可以预先处理出互质的数对。

我们可以维护一个数组表示「某一时刻离每个节点值距离最近的互斥节点值的节点号」。在遍历到某一节点时，我们可以将预先处理好的与该节点值互质的节点值刷新为当前节点号。

AC代码

class Solution {
    fun getCoprimes(nums: IntArray, edges: Array<IntArray>): IntArray {
    
        //预处理
        val arr = Array<ArrayList<Int>>(51) { arrayListOf() }
        for (i in 1..50) for (j in 1..50) if (gcd(i, j) == 1) arr[i].add(j)
        
        //构造树
        val map = HashMap<Int, ArrayList<Int>>()
        edges.forEach {
            map[it[0]] = (map[it[0]] ?: arrayListOf()).apply { add(it[1]) }
            map[it[1]] = (map[it[1]] ?: arrayListOf()).apply { add(it[0]) }
        }
        
        val res = IntArray(nums.size)
        
        //回溯
        val c = IntArray(51) { -1 }
        fun dfs(cur: Int, parent: Int) {
            res[cur] = c[nums[cur]]
            arr[nums[cur]].forEach { c[it] = cur }   // 刷新状态
            map[cur]?.forEach {
                if (it != parent) {
                    val t = IntArray(51)
                    for (i in t.indices) t[i] = c[i]  // 保存案发现场
                    dfs(it, cur)
                    for (i in t.indices) c[i] = t[i]  // 退回上一步
                }
            }   
        }
        dfs(0, -1)
        return res
    }

    fun gcd(a: Int, b: Int): Int {
        if (b == 0) return a
        return gcd(b, a % b)
    }
}