Minimum Time to Collect All Apples in a Tree LeetCode solution in C++

1
class DSU {
2
 private:
3
  vector<int> parent, rank;
4

5
 public:
6
  DSU(int n) {
7
    rank.resize(n, 1);
8
    parent.resize(n);
9
    iota(begin(parent), end(parent), 0);
10
  }
11

12
  int find_parent(int node) {
13
    if (node == parent[node]) return node;
14
    return parent[node] = find_parent(parent[node]);
15
  }
16
  void Union(int u, int v) {
17
    int U = find_parent(u), V = find_parent(v);
18
    if (U == V) return;
19
    if (rank[U] < rank[V]) swap(U, V);
20
    rank[U] += rank[V];
21
    parent[V] = U;
22
  }
23
  int getRank(int node) {
24
    return rank[parent[node]];
25
  }
26
};
27

28
class Solution {
29
 public:
30
  vector<bool> dp;
31
  vector<bool> hasApple;
32
  vector<int> adj[100001];
33
  vector<int> visited;
34

35
  void dfs(int src) {
36
    visited[src] = 1;
37
    for (auto &nbr : adj[src]) {
38
      if (!visited[nbr]) {
39
        dfs(nbr);
40
        dp[src] = dp[src] or dp[nbr];
41
      }
42
    }
43
  }
44

45
  int minTime(int n, vector<vector<int>> &edges, vector<bool> &hasApple) {
46
    DSU dsu(n);
47
    dp = hasApple;
48
    visited.resize(n, 0);
49
    this->hasApple = hasApple;
50

51
    for (auto &x : edges) adj[x[0]].push_back(x[1]), adj[x[1]].push_back(x[0]);
52
    dfs(0);
53

54
    int start = -1;
55
    for (int i = 0; i < n; ++i) {
56
      if (!dp[i]) continue;
57
      if (start == -1) {
58
        start = i;
59
        continue;
60
      }
61
      dsu.Union(start, i);
62
    }
63

64
    return (dsu.getRank(0) - 1) * 2;
65
  }
66
};