https://leetcode.com/problems/tree-of-coprimes/ from Google

Thinking process:

First observation is that the number of nodes can be up to 10^5, which means we can only use O(NlogN) algoritm in the worst case. So if the tree is balanced, then just compare all the ancestors will be O(NlogN), but here it doesn’t state that the tree is balanced, so comparing all ancestors doesn’t work.

Next notice that the value of nums is bounded at 50, which means we only have at most 50 possible ancestor values. This seems promising. We can use an array of size 50 to dedup the ancestors with the same value. The only problem is how can we know which ancestor is the nearest one? We can solve this by using the depth of the ancestor. Here is an example.

If we’re at node index 2, we should have two arrays: one is the parentToIndex, another one is the indexToDepth. The parentToIndex is a map of the ancestor’s value to a stack of ancestor’s indexes. As you can see, our ancestors’ indexes are 0, 3, and 1. Node 1 and 3 has the same value 3, so they are in the parentToIndex[3]’s stack. Node 0’s value is 4, so it’s in the parentToIndex[4]’s stack.

With the parentToIndex, to find the coprime, we only need to compare at most 50 times as all nodes in the same stack are having the same value. After we find the coprime values, we now need to use indexToDepth to find the deepest index as that will be the node cloest to us.

So in this example, we know that node 0 and node 1 are our coprimes (ignore node 3 as 1 is closer to us). Then we check the depth of node 0 and node 1, and know that 1 is the answer we want.

Here is the code.

 1class Solution {
 2public:
 3    vector<int> ans;
 4    vector<int> ns;
 5    vector<vector<int>> parentToIndex;
 6    vector<int> indexToDepth;
 7    void dfs(int from, int cur, int depth, vector<vector<int>>& es) {
 8        vector<int> cand;
 9        indexToDepth[cur] = depth;
10        depth++;
11        for (int i=1; i<=50; i++) {
12            if (parentToIndex[i].size() and gcd(ns[cur], i) == 1) {
13                cand.push_back(parentToIndex[i].back());
14            }
15        }
16        if (cand.empty()) ans[cur] = -1;
17        else {
18            int t = cand[0];
19            for (auto c : cand) {
20                if (indexToDepth[t] < indexToDepth[c])
21                    t = c;
22            }
23            ans[cur] = t;
24        }
25        parentToIndex[ns[cur]].push_back(cur);
26        for (auto& nxt : es[cur]) {
27            if (nxt != from) {
28                dfs(cur, nxt, depth, es);
29            }
30        }
31        parentToIndex[ns[cur]].pop_back();
32    }
33    
34    vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
35        int N = nums.size();
36        ns = nums;
37        vector<vector<int>> es(N);
38        parentToIndex.resize(51);
39        indexToDepth.resize(N);
40        ans.resize(N);
41        for (auto& e : edges) {
42            es[e[0]].push_back(e[1]);
43            es[e[1]].push_back(e[0]);
44        }
45        dfs(-1, 0, 0, es);
46        return ans;
47    }
48};