LeetCode in Kotlin

1766. Tree of Coprimes

Hard

There is a tree (i.e., a connected, undirected graph that has no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. Each node has a value associated with it, and the root of the tree is node 0.

To represent this tree, you are given an integer array nums and a 2D array edges. Each nums[i] represents the ith node’s value, and each edges[j] = [uj, vj] represents an edge between nodes uj and vj in the tree.

Two values x and y are coprime if gcd(x, y) == 1 where gcd(x, y) is the greatest common divisor of x and y.

An ancestor of a node i is any other node on the shortest path from node i to the root. A node is not considered an ancestor of itself.

Return an array ans of size n, where ans[i] is the closest ancestor to node i such that nums[i] and nums[ans[i]] are coprime, or -1 if there is no such ancestor.

Example 1:

Input: nums = [2,3,3,2], edges = [[0,1],[1,2],[1,3]]

Output: [-1,0,0,1]

Explanation: In the above figure, each node’s value is in parentheses.

Example 2:

Input: nums = [5,6,10,2,3,6,15], edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]

Output: [-1,0,-1,0,0,0,-1]

Constraints:

Solution

@Suppress("kotlin:S107")
class Solution {
    private fun dfs(
        v2n: IntArray,
        v2d: IntArray,
        depth: Int,
        parent: Int,
        node: Int,
        ans: IntArray,
        nums: IntArray,
        neighbors: Array<ArrayList<Int>>
    ) {
        var d = Int.MIN_VALUE
        var n = -1
        val v = nums[node]
        for (i in 1..50) {
            if (v2n[i] != -1 && v2d[i] > d && gcd(i, v) == 1) {
                d = v2d[i]
                n = v2n[i]
            }
        }
        ans[node] = n
        val v2NOld = v2n[v]
        val v2DOld = v2d[v]
        v2n[v] = node
        v2d[v] = depth
        for (child in neighbors[node]) {
            if (child == parent) {
                continue
            }
            dfs(v2n, v2d, depth + 1, node, child, ans, nums, neighbors)
        }
        v2n[v] = v2NOld
        v2d[v] = v2DOld
    }

    private fun gcd(x: Int, y: Int): Int {
        return if (x == 0) y else gcd(y % x, x)
    }

    fun getCoprimes(nums: IntArray, edges: Array<IntArray>): IntArray {
        val n = nums.size
        val neighbors: Array<ArrayList<Int>> = Array(n) { ArrayList() }
        for (i in 0 until n) {
            neighbors[i] = ArrayList()
        }
        for (edge in edges) {
            neighbors[edge[0]].add(edge[1])
            neighbors[edge[1]].add(edge[0])
        }
        val ans = IntArray(n)
        val v2n = IntArray(51)
        val v2d = IntArray(51)
        v2n.fill(-1)
        dfs(v2n, v2d, 0, -1, 0, ans, nums, neighbors)
        return ans
    }
}