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. - Node 0 has no coprime ancestors. - Node 1 has only one ancestor, node 0. Their values are coprime (gcd(2,3) == 1). - Node 2 has two ancestors, nodes 1 and 0. Node 1's value is not coprime (gcd(3,3) == 3), but node 0's value is (gcd(2,3) == 1), so node 0 is the closest valid ancestor. - Node 3 has two ancestors, nodes 1 and 0. It is coprime with node 1 (gcd(3,2) == 1), so node 1 is its closest valid ancestor.
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:
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
自己写的TLE的代码
class Solution {
List[] gcdSet = new List[51];
//记录所有的父节点的val和节点编号
List<int[]> path = new ArrayList();
int[] ans;
int N = 100000;
Map<Integer, Set<Integer>> g = new HashMap(N);
public int[] getCoprimes(int[] nums, int[][] edges) {
int n = nums.length, m = edges.length;
for(int i = 0; i < m; i++){
int a = edges[i][0], b = edges[i][1];
g.putIfAbsent(a, new HashSet());
g.putIfAbsent(b, new HashSet());
g.get(a).add(b); g.get(b).add(a);
}
ans = new int[n];
Arrays.fill(ans, -1);
for(int i = 1; i <= 50; i++){
gcdSet[i] = new ArrayList();
for(int j = 1; j <= 50; j++){
if(gcd(i, j) == 1) gcdSet[i].add(j);
}
}
dfs(0, -1, nums);
return ans;
}
public void dfs(int cur, int father, int[] nums){
if(g.get(cur) == null) return;
int curVal = nums[cur];
for(int i = path.size() - 1; i >= 0; i--){
if(gcdSet[nums[cur]].contains(path.get(i)[0])){
ans[cur] = path.get(i)[1];
break;
}
}
path.add(new int[]{curVal, cur});
for(int ne: g.get(cur)){
if(ne == father) continue;
dfs(ne, cur, nums);
}
path.remove(path.size() - 1);
}
public int gcd(int a, int b){
return b == 0 ? a: gcd(b, a % b);
}
}
TLE就是因为,没有办法O(1)判断最近的互质数
加入pos记录值的节点和depth记录节点深度
class Solution {
int N = 100000;
int[] pos = new int[51]; // 每个值对应的点的位置, 因为只需要记录最下面的点,所以开数组即可
int[] ans, depth; // 存每个点所在的深度
List<Integer>[] gcdList = new List[51];
Map<Integer, Set<Integer>> g = new HashMap(N);
public int[] getCoprimes(int[] nums, int[][] edges) {
int n = nums.length, m = edges.length;
init();
for(int i = 0; i < m; i++){
int a = edges[i][0], b = edges[i][1];
g.putIfAbsent(a, new HashSet());
g.putIfAbsent(b, new HashSet());
g.get(a).add(b); g.get(b).add(a);
}
ans = new int[n];
depth = new int[n];
Arrays.fill(ans, -1);
Arrays.fill(pos, -1);
dfs(nums, 0, -1, 0);
return ans;
}
public void dfs(int[] nums, int u, int f, int idx){
if(g.get(u) == null) return;
int val = nums[u];
for(int j: gcdList[val]){
if(ans[u] == -1 || (pos[j] != -1 && depth[pos[j]] > depth[ans[u]])){
ans[u] = pos[j];
}
}
int t = pos[val];
depth[u] = idx;
pos[val] = u;
for(int ne: g.get(u)){
if(ne == f) continue;
dfs(nums, ne, u, idx + 1);
}
pos[val] = t;
}
public int gcd(int a, int b){
return b == 0 ? a: gcd(b, a % b);
}
public void init(){
for(int i = 1; i <= 50; i++){
gcdList[i] = new ArrayList();
for(int j = 1; j <= 50; j++){
if(gcd(i, j) == 1) gcdList[i].add(j);
}
}
}
}