LeetCode 2581. Count Number of Possible Root Nodes

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
class Solution {
public:
int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
vector<vector<int>> g(edges.size() + 1);
for (auto &e : edges) {
g[e[0]].push_back(e[1]);
g[e[1]].push_back(e[0]);
}
unordered_set<long> s;
for (auto &gs : guesses) s.insert((long) gs[0] << 32 | gs[1]);
int ans = 0, cnt0 = 0;
function<void(int, int)> dfs = [&](int x, int f) {
for (int y : g[x]) {
if (y != f) {
cnt0 += s.count((long) x << 32 | y);
dfs(y, x);
}
}
};
dfs(0, -1);
function<void(int, int, int)> reroot = [&](int x, int f, int cnt) {
if (cnt >= k) ++ans;
for (int y : g[x]) {
if (y != f) {
reroot(y, x, cnt
- s.count((long) x << 32 | y)
+ s.count((long) y << 32 | x));
}
}
};
reroot(0, -1, cnt0);
return ans;
}
};