DFS - 开数组记录路径 - Tree of Coprimes(hard) - Biweekly Contest 46

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);
            }
        }
    }
}