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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
 class Node { public: Node(int i) : index(i), parent_index(1), is_prisoner(false), has_escape_leaf(true), has_escape_root(false) {} int index, parent_index; bool is_prisoner, has_escape_leaf, has_escape_root; vector<int> neis; };
using MapType = unordered_map<int, Node*>;
void dfs(MapType& graph, int root, int& res, unordered_map<int, bool>& visited) { visited[root] = true; Node* cur = graph[root]; for_each(cur>neis.begin(), cur>neis.end(), [&graph, &visited, &res, root](int nei){ if (visited.find(nei) == visited.end()) { graph[nei]>parent_index = root; dfs(graph, nei, res, visited); } }); if (cur>parent_index != 1 && cur>neis.size() == 1) return; int subtree_num = 0, escape_leaf = 0, escape_root = 0; for_each(cur>neis.begin(), cur>neis.end(), [&](int nei){ if (nei == cur>parent_index) return; subtree_num++; if (graph[nei]>has_escape_leaf) escape_leaf++; if (graph[nei]>has_escape_root) escape_root++; }); if (cur>is_prisoner) { res += escape_leaf; cur>has_escape_leaf = false; cur>has_escape_root = true; } else { if (escape_leaf == subtree_num) { } else if (escape_leaf == 0) { cur>has_escape_leaf = false; if (escape_root) cur>has_escape_root = true; } else { if (escape_root) { res++; cur>has_escape_root = false; cur>has_escape_leaf = false; } else { cur>has_escape_leaf = true; cur>has_escape_root = false; } } } }
int solution(vector<int> &A, vector<int> &B, vector<int> &C) { if (C.empty()) return 0; MapType graph; unordered_map<int, bool> visited; int n = static_cast<int>(A.size()), m = static_cast<int>(C.size()); for (int i = 0; i < n; i++) { if (graph.find(A[i]) == graph.end()) graph.insert(MapType::value_type(A[i], new Node(A[i]))); if (graph.find(B[i]) == graph.end()) graph.insert(MapType::value_type(B[i], new Node(B[i]))); graph[A[i]]>neis.push_back(B[i]); graph[B[i]]>neis.push_back(A[i]); } for (int i = 0; i < m; i++) { int prionser = C[i]; if (graph[prionser]>neis.size() == 1) return 1; graph[prionser]>is_prisoner = true; } int res = 0; dfs(graph, 0, res, visited); if (graph[0]>has_escape_root && graph[0]>neis.size() == 1) res++; return res; }
