3249.cpp (1467B)
1 class Solution { 2 public: 3 int countGoodNodes(const vector<vector<int>> &edges) const { 4 static int count[100001]; 5 const int n = size(edges) + 1; 6 vector<vector<int>> adj(n); 7 stack<pair<int, int>> st; 8 int res = 0; 9 10 for (const auto &edge : edges) { 11 adj[edge[0]].push_back(edge[1]); 12 adj[edge[1]].push_back(edge[0]); 13 } 14 15 st.emplace(0, -1); 16 memset(count, 0x00, sizeof(count)); 17 while (!st.empty()) { 18 if (st.top().first != -1) { 19 const auto [root, parent] = st.top(); 20 st.emplace(-1, -1); 21 22 for (const int next : adj[root]) { 23 if (next == parent) continue; 24 st.emplace(next, root); 25 } 26 27 continue; 28 } 29 30 st.pop(); 31 const auto [root, parent] = st.top(); 32 st.pop(); 33 34 int cnt = 1; 35 int goal = -1; 36 bool good = true; 37 38 for (int i = 0; i < size(adj[root]); i++) { 39 const int next = adj[root][i]; 40 if (next == parent) continue; 41 if (goal == -1) 42 goal = count[next]; 43 else if (count[next] != goal) 44 good = false; 45 cnt += count[next]; 46 } 47 48 if (good) res++; 49 50 count[root] = cnt; 51 } 52 53 return res; 54 } 55 };