HourRank 20: Birjik and Nicole's Tree Game
https://www.hackerrank.com/contests/hourrank-20/challenges/birjik-and-nicoles-tree-game
解法
黒頂点を post-order でソートする。この列を左から順に見ていき、隣接する 2 点の LCA を列に挿入することを繰り返す。これだけで LCA の閉包(?)が作れる。
HL 分解は LCA を求める部分にしか使っていない。
#include <iostream> #include <algorithm> #include <string> #include <vector> #include <functional> #include <queue> #include <stack> #include <map> #include <set> struct HLDecomp { using Tree = std::vector<std::vector<int>>; const Tree &tree; std::vector<int> parent; std::vector<int> head; std::vector<int> vid; HLDecomp(const Tree &tree) : tree(tree) { const int n = tree.size(); const int root = 0; std::stack<std::pair<int, int>> stack; stack.emplace(root, 0); parent.assign(n, -1); head.assign(n, -1); vid.assign(n, -1); std::vector<int> heavy(n, -1); std::vector<int> size(n, 1); while (!stack.empty()) { const int u = stack.top().first; const int i = stack.top().second; if (i < tree[u].size()) { stack.top().second++; const int v = tree[u][i]; if (v != parent[u]) { parent[v] = u; stack.emplace(v, 0); } } else { stack.pop(); int max = 0; for (int v : tree[u]) { if (v != parent[u]) { size[u] += size[v]; if (max < size[v]) { max = size[v]; heavy[u] = v; } } } } } std::queue<int> queue; queue.push(0); int now = 0; while (!queue.empty()) { const int h = queue.front(); queue.pop(); for (int i = h; i != -1; i = heavy[i]) { head[i] = h; vid[i] = now++; for (int j : tree[i]) { if (j != parent[i] && j != heavy[i]) { queue.push(j); } } } } } template<typename T> void foreach(int u, int v, T func) { while (true) { if (vid[u] > vid[v]) { std::swap(u, v); } if (head[u] != head[v]) { func(vid[head[v]], vid[v]); v = parent[head[v]]; } else { func(vid[u], vid[v]); break; } } } int lca(int u, int v) { while (true) { if (vid[u] > vid[v]) { std::swap(u, v); } if (head[u] != head[v]) { v = parent[head[v]]; } else { return u; } } } }; int main() { int n; std::cin >> n; std::vector<std::vector<int>> tree(n); for (int i = 0; i < n - 1; i++) { int u, v; scanf("%d %d", &u, &v); u--; v--; tree[u].push_back(v); tree[v].push_back(u); } HLDecomp hld(tree); int q; std::cin >> q; std::vector<int> post(n); std::vector<int> ord; std::vector<int> depth(n); int k = 0; std::function<void(int, int)> dfs = [&](int u, int p) { for (int v : tree[u]) { if (v != p) { depth[v] = depth[u] + 1; dfs(v, u); } } post[u] = k++; ord.push_back(u); }; dfs(0, -1); std::vector<int> imos(n); while (q--) { int k; scanf("%d", &k); std::set<int> q; q.insert(post[0]); for (int i = 0; i < k; i++) { int t; scanf("%d", &t); t--; q.insert(post[t]); imos[t]++; } std::vector<int> a; std::map<int, int> parent; while (!q.empty()) { int x = ord[*q.begin()]; q.erase(q.begin()); a.push_back(x); if (!q.empty()) { int y = ord[*q.begin()]; int z = hld.lca(x, y); if (z != y) { parent[y] = z; q.insert(post[z]); } parent[x] = z; } } std::vector<int> ans(k + 1); for (int i : a) { if (i != 0) { int p = parent[i]; imos[p] += imos[i]; ans[imos[i]] += depth[i] - depth[p]; } else { ans[imos[i]]++; } } ans[0] = n; for (int i = 1; i < k + 1; i++) { ans[0] -= ans[i]; } for (int i = 0; i < k + 1; i++) { printf("%d ", ans[i]); } printf("\n"); for (int i : a) { imos[i] = 0; } } }