leetcode

0 class Solution {

1 public:

2 int countGoodNodes(const vector<vector<int>> &edges) const {

3 static int count[100001];

4 const int n = size(edges) + 1;

5 vector<vector<int>> adj(n);

6 stack<pair<int, int>> st;

7 int res = 0;

8

9 for (const auto &edge : edges) {

10 adj[edge[0]].push_back(edge[1]);

11 adj[edge[1]].push_back(edge[0]);

12 }

13

14 st.emplace(0, -1);

15 memset(count, 0x00, sizeof(count));

16 while (!st.empty()) {

17 if (st.top().first != -1) {

18 const auto [root, parent] = st.top();

19 st.emplace(-1, -1);

20

21 for (const int next : adj[root]) {

22 if (next == parent) continue;

23 st.emplace(next, root);

24 }

25

26 continue;

27 }

28

29 st.pop();

30 const auto [root, parent] = st.top();

31 st.pop();

32

33 int cnt = 1;

34 int goal = -1;

35 bool good = true;

36

37 for (int i = 0; i < size(adj[root]); i++) {

38 const int next = adj[root][i];

39 if (next == parent) continue;

40 if (goal == -1)

41 goal = count[next];

42 else if (count[next] != goal)

43 good = false;

44 cnt += count[next];

45 }

46

47 if (good) res++;

48

49 count[root] = cnt;

50 }

51

52 return res;

53 }

54 };