AOJ 2170 Marked Ancestor
http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2170&lang=jp
蟻本にUnion-Find木の練習問題として掲載されている問題です。愚直解があっさり通ってしまうのが悲しいところですが、想定解法はわりと難しいと思います。
概要は以下のとおり。
木が与えられます。この木に対し、Q個のクエリがあるのでそれに答えて下さい。クエリは2種類。
M : あるノードを赤く(色指定はなかったかも)塗る
Q : あるノードについて、そのノード自身と、親、そのまた親……の中から、マークされているノードで一番近いものを選ぶ。
選んだノードの和を出力して下さい。
というもの。
まず頭悪そうな愚直解を投げたところ0.22秒で通ってしまいました。↓が短く書き直した愚直解です。
int T[100000]; bool marked[100000]; int dfs(int v){ return marked[v] ? v : dfs(T[v]);} int main(){ cin.tie(0); ios::sync_with_stdio(false); int N, Q, _t; char _c; while(cin >> N >> Q && N){ memset(marked, 0, sizeof(marked)); REP(i, 1, N){ cin >> _t; _t--; T[i] = _t; } marked[0] = true; ll sum = 0; rep(i, Q){ cin >> _c >> _t; _t--; if(_c == 'Q') sum += dfs(_t) + 1; else marked[_t] = true; } cout << sum << endl; } return 0; }
ではどうやってこれを速くするか。
木構造なのでDFSはわりと効率がいいのですが、いちいち辿っていたのではやはり無駄っぽいです。あるノードvに対し、Q(v)が0であるときを考えます。このときvが塗られたら、Q(v)はその後ずっとvになります。これを効率よく再現する……のは大変そうなので、クエリを逆回しすることを考えましょう。すると、vがもとの色に塗り直される、という操作は、vを根とする部分木が0を根とする部分木に併合される、という操作だと解釈できます。なので、Union-Find木を使って
1:M(v)で指定された全部の頂点を塗る
2 : 塗られていない頂点については、その親と併合する。この操作により、すべての塗られた頂点について、自らが根となる部分木ができる。
3:クエリを逆回しし、Q(v)ではvの根をsumに加え、M(v)ではvとvの親を併合する。
という操作により問題が解けます。0.10秒程度でした。
実装上の注意点として、Unionしたとき、必ずもとの木で親となるものが根になるようにします。このため、UnionFind木中のparをすべてマイナス1で初期化し、親が常に負の値を取るような実装にします。また、呼び出すときは常に(親、子)の順で呼び出します。
struct UnionFind{ vector <int> par; UnionFind(int N) : par(N, -1){} void unite(int x, int y){ x = root(x); y = root(y); if(x != y){ par[x] += par[y]; par[y] = x; } } int root(int x){ return par[x] < 0 ? x : par[x] = root(par[x]);} bool same(int x, int y){ return root(x) == root(y);} }; typedef pair<bool, int> pbt;; bool marked[100000]; int p[100000]; int main(){ cin.tie(0); ios::sync_with_stdio(false); int N, Q, _t; char _c; while(ifs >> N >> Q && N){ memset(marked, 0, sizeof(marked)); marked[0] = true; UnionFind uf(N); vector<pbt> query; REP(i, 1, N){ ifs >> p[i]; --p[i]; } rep(i, Q){ ifs >> _c >>_t; --_t; if(_c == 'Q') query.push_back(pbt(1, _t)); else if(!marked[_t]){ query.push_back(pbt(0, _t)); marked[_t] = true; } } rep(i, N) if(!marked[i]) uf.unite(p[i], i); ll sum = 0LL; for(int i = query.size() - 1; i >= 0; --i){ int v = query[i].second; if(query[i].first) sum += uf.root(v) + 1; else uf.unite(p[v], v); } cout << sum << endl; } return 0; }
想定解のほうは解説を見てもよくわからず少し苦労しました。何時間もかかったわけではないですが……。