Skip to content

1766. Tree of Coprimes 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class Solution {
 public:
  vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
    vector<int> ans(nums.size(), -1);
    vector<vector<int>> tree(nums.size());
    // stacks[i] := (node, depth)s of nodes with value i
    vector<stack<pair<int, int>>> stacks(kMax + 1);

    for (const vector<int> edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].push_back(v);
      tree[v].push_back(u);
    }

    dfs(tree, 0, /*prev=*/-1, /*depth=*/0, nums, stacks, ans);
    return ans;
  }

 private:
  static constexpr int kMax = 50;

  void dfs(const vector<vector<int>>& tree, int u, int prev, int depth,
           const vector<int>& nums, vector<stack<pair<int, int>>>& stacks,
           vector<int>& ans) {
    ans[u] = getAncestor(u, stacks, nums);
    stacks[nums[u]].push({u, depth});

    for (const int v : tree[u])
      if (v != prev)
        dfs(tree, v, u, depth + 1, nums, stacks, ans);

    stacks[nums[u]].pop();
  }

  int getAncestor(int u, const vector<stack<pair<int, int>>>& stacks,
                  const vector<int>& nums) {
    int maxNode = -1;
    int maxDepth = -1;
    for (int i = 1; i <= kMax; ++i)
      if (!stacks[i].empty() && stacks[i].top().second > maxDepth &&
          __gcd(nums[u], i) == 1) {
        maxNode = stacks[i].top().first;
        maxDepth = stacks[i].top().second;
      }
    return maxNode;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class Solution {
  public int[] getCoprimes(int[] nums, int[][] edges) {
    int[] ans = new int[nums.length];
    Arrays.fill(ans, -1);
    List<Integer>[] tree = new List[nums.length];
    // stacks[i] := (node, depth)s of nodes with value i
    Deque<Pair<Integer, Integer>>[] stacks = new Deque[kMax + 1];

    for (int i = 0; i < nums.length; ++i)
      tree[i] = new ArrayList<>();

    for (int i = 1; i <= kMax; ++i)
      stacks[i] = new ArrayDeque<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      tree[u].add(v);
      tree[v].add(u);
    }

    dfs(tree, 0, /*prev=*/-1, /*depth=*/0, nums, stacks, ans);
    return ans;
  }

  private static final int kMax = 50;

  private void dfs(List<Integer>[] tree, int u, int prev, int depth, int[] nums,
                   Deque<Pair<Integer, Integer>>[] stacks, int[] ans) {
    ans[u] = getAncestor(u, stacks, nums);
    stacks[nums[u]].push(new Pair<>(u, depth));

    for (final int v : tree[u])
      if (v != prev)
        dfs(tree, v, u, depth + 1, nums, stacks, ans);

    stacks[nums[u]].pop();
  }

  private int getAncestor(int u, Deque<Pair<Integer, Integer>>[] stacks, int[] nums) {
    int maxNode = -1;
    int maxDepth = -1;
    for (int i = 1; i <= kMax; ++i)
      if (!stacks[i].isEmpty() && stacks[i].peek().getValue() > maxDepth && gcd(nums[u], i) == 1) {
        maxNode = stacks[i].peek().getKey();
        maxDepth = stacks[i].peek().getValue();
      }
    return maxNode;
  }

  private int gcd(int a, int b) {
    return b == 0 ? a : gcd(b, a % b);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution:
  def getCoprimes(self, nums: list[int], edges: list[list[int]]) -> list[int]:
    kMax = 50
    ans = [-1] * len(nums)
    tree = [[] for _ in range(len(nums))]
    # stacks[i] := (node, depth)s of nodes with value i
    stacks = [[] for _ in range(kMax + 1)]

    for u, v in edges:
      tree[u].append(v)
      tree[v].append(u)

    def getAncestor(u: int) -> int:
      maxNode = -1
      maxDepth = -1
      for i, stack in enumerate(stacks):
        if stack and stack[-1][1] > maxDepth and math.gcd(nums[u], i) == 1:
          maxNode, maxDepth = stack[-1]
      return maxNode

    def dfs(u: int, prev: int, depth: int) -> int:
      ans[u] = getAncestor(u)
      stacks[nums[u]].append((u, depth))

      for v in tree[u]:
        if v != prev:
          dfs(v, u, depth + 1)

      stacks[nums[u]].pop()

    dfs(0, -1, 0)
    return ans