golang 解法, 执行用时: 476 ms, 内存消耗: 59.7 MB, 提交时间: 2023-10-25 23:26:22

var prime = make([][]bool, 0)

func getCoprimes(nums []int, edges [][]int) []int {
	var edge = make([][]int, len(nums))

	for i := range edge {
		edge[i] = make([]int, 0)
	}
	for _, e := range edges {
		a, b := e[0], e[1]
		edge[a] = append(edge[a], b)
		edge[b] = append(edge[b], a)
	}
	if len(prime) == 0 {
		for i := 0; i < 51; i++ {
			prime = append(prime, make([]bool, 51))
		}
		for i := 0; i < 51; i++ {
			for j := i; j < 51; j++ {
				prime[j][i] = greatestCommonDivisor(i, j) == 1
				prime[i][j] = prime[j][i]
			}
		}
	}

	var set = make(map[int]bool, 0)
	for _, n := range nums {
		set[n] = true
	}
	var visited = make(map[int]bool, 0)
	var currents = make([][][]int, 51)
	for i := 0; i < 51; i++ {
		currents[i] = make([][]int, 0)
	}

	var results = make([]int, len(nums))
	var dfs func(node, depth int)
	dfs = func(node, depth int) {
		if visited[node] {
			return
		}
		visited[node] = true
		var value = nums[node]
		var ans = []int{-1, -1}
		for i := range set {
			if len(currents[i]) > 0 && prime[value][i] {
					var get = currents[i][len(currents[i])-1]
				if get[1] > ans[1] {
					ans = get
				}
			}
		}
		results[node] = ans[0]

		for _, child := range edge[node] {
			currents[value] = append(currents[value], []int{node, depth})

			dfs(child, depth+1)
			currents[value] = currents[value][0 : len(currents[value])-1]
		}
	}
	dfs(0, 0)

	return results
}

func greatestCommonDivisor(a int, b int) int {
	for b != 0 {
		return greatestCommonDivisor(b, a%b)
	}
	return a
}

python3 解法, 执行用时: 2084 ms, 内存消耗: 117.9 MB, 提交时间: 2023-10-25 23:25:54

class Solution:
    def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:
        graph = collections.defaultdict(list)

        for edge in edges:
            i, j = edge
            graph[i].append(j)
            graph[j].append(i)

        stacks = [[] for _ in range(51)]
        res = [-1 for i in range(len(nums))]

        def dfs(node, parent, depth):
            candidates = []
            for i in range(1, len(stacks)):
                if len(stacks[i]) > 0 and math.gcd(i, nums[node]) == 1:
                    candidates.append(stacks[i][-1])

            maxDepth = -1
            resNode = -1
            for candidate in candidates:
                d, n = candidate
                if d > maxDepth:
                    maxDepth = d
                    resNode = n
            res[node] = resNode

            stacks[nums[node]].append((depth, node))

            children = graph[node]
            for child in children:
                if child != parent:
                    dfs(child, node, depth + 1)
            stacks[nums[node]].pop()

        dfs(0, None, 0)
        return res

cpp 解法, 执行用时: 556 ms, 内存消耗: 172.6 MB, 提交时间: 2023-10-25 23:24:58

class Solution {
public:
    vector<vector<int>> G;
    stack<pair<int,int>> lasts[55];
    vector<int> res;
    void dfs(int node, int pre, int level, vector<int>& a) {
        int re = -1, lev = -1;
        for(int i = 1; i <= 50; ++i) {
            if(lasts[i].size() && lasts[i].top().first > lev && __gcd(i, a[node]) == 1) {
                re = lasts[i].top().second;
                lev = lasts[i].top().first;
            }
        }
        res[node] = re;
        for(int ne : G[node]) {
            if(ne != pre) {
                lasts[a[node]].push({level, node});
                dfs(ne, node, level + 1, a);
                lasts[a[node]].pop();
            }
        }
    }
    vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
        int n = nums.size();
        G.resize(n);
        for(auto& e : edges) {
            G[e[0]].push_back(e[1]);
            G[e[1]].push_back(e[0]);
        }
        res.resize(n);
        dfs(0, -1, 0, nums);
        return res;
    }
};