Codeforces Round #292 (Div. 1) D. Drazil and Morning Exercise
問題概要
木 T が与えられる.また,クエリが q 回投げられる.それぞれのクエリは L を投げてくる.
木 T の頂点 i に対して,i からそこから最も遠い葉までの距離を D(i) とする.
各クエリに対し,T の部分木 S で,かつ S の任意の二頂点 i, j が |D(i) - D(j)| <= L となるようなものの中で,最も大きいものの大きさを求めよ.
D(i) はあくまでも T での距離であって,S での距離ではないことに注意.
・制約
1 <= n <= 10^5
1 <= q <= 50
1 <= L <= 10^11
解法
全方位木DP,Union-Find木で解ける.
T においてもっとも遠い葉までの距離を,全方位木DPで求める.
そうして得られた距離 d[i] を,降順にソートする.
あとはソートして得られた距離の列上で,前からしゃくとりする.最大が x なら x - L までは含んで良い.
しゃくとりで新たに頂点を追加する場合,その頂点の隣接点で,今現在追加されている頂点があるなら,そこと union-find でマージする.
しゃくとりで左端を縮めるときは,union-find のサイズを1つ小さくするだけでよい.
これは,d を降順でソートしているからである.もし昇順にしてしまうと,サイズを1減らすだけではおかしな値になる.
(イメージとしては,葉の方から根の方にどんどん集合が集まってくる感じ.根からスタートすると,途中で分岐したときに集合を完全にわけないとだめになって,サイズを 1 減らすだけでは対応できない.葉の方からスタートすると,途中でくっつくことがあっても別れることはない.)
ソースコード
#include <bits/stdc++.h> using namespace std; using ll = long long; using pll = pair<ll, ll>; struct edge { int to; ll cost; }; using edges = vector<edge>; using graph = vector<edges>; class union_find { public: union_find(int n) : par_(n, -1) {} int root(int x) { return par_[x] < 0 ? x : par_[x] = root(par_[x]); } bool unite(int x, int y) { x = root(x); y = root(y); if(x == y) { return false; } else { if(par_[x] < par_[y]) { par_[x] += par_[y]; par_[y] = x; } else { par_[y] += par_[x]; par_[x] = y; } return true; } } bool same(int x, int y) { return root(x) == root(y); } int size(int x) { return -par_[root(x)]; } void dec(int x) { par_[root(x)] += 1; } private: std::vector<int> par_; }; void dfs(int v, int prev, graph const& g, vector<ll>& d) { d[v] = 0; for(auto& e : g[v]) { if(e.to == prev) { continue; } dfs(e.to, v, g, d); d[v] = max(d[v], d[e.to] + e.cost); } } void dfs2(int v, int prev, ll par_d, graph const& g, vector<ll> const& d, vector<pll>& res) { res[v].second = v; vector<pll> child = {{0, -1}}; for(auto& e : g[v]) { if(e.to == prev) { child.emplace_back(par_d + e.cost, e.to); } else { child.emplace_back(d[e.to] + e.cost, e.to); } } sort(child.rbegin(), child.rend()); res[v].first = child[0].first; for(auto& e : g[v]) { if(e.to == prev) { continue; } if(e.to == child[0].second) { dfs2(e.to, v, child[1].first, g, d, res); } else { dfs2(e.to, v, child[0].first, g, d, res); } } } int main() { int n; cin >> n; graph g(n); for(int i = 0; i < n - 1; ++i) { int x, y; ll v; cin >> x >> y >> v; x--; y--; g[x].push_back(edge{y, v}); g[y].push_back(edge{x, v}); } vector<ll> tmp_d(n); vector<pll> d(n); dfs(0, -1, g, tmp_d); dfs2(0, -1, 0, g, tmp_d, d); sort(d.rbegin(), d.rend()); int q; cin >> q; while(q--) { ll L; cin >> L; int l = 0, r = 0; vector<bool> used(n); union_find uf(n); int res = 0; while(r < n) { while(r < n && d[l].first <= d[r].first + L) { for(auto& e : g[d[r].second]) { if(used[e.to]) { uf.unite(e.to, d[r].second); } } used[d[r].second] = true; r++; } res = max(res, uf.size(d[l].second)); uf.dec(d[l].second); used[d[l].second] = false; l++; } cout << res << endl; } }
感想
オーバーフローに気づかなくで大量に時間を溶かしてしまった.
UnionFind ってこういう使い方もあるんだなあと勉強になった.