読者です 読者をやめる 読者になる 読者になる

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;
}

 想定解のほうは解説を見てもよくわからず少し苦労しました。何時間もかかったわけではないですが……。