AOJ - 2170 Marked Ancestor

問題を解いた

問題概要

木が与えられる。以下の２つのクエリを処理し、出力された番号の合計を求める。

頂点aをマークする
頂点aから最も近いマークされた祖先の番号を出力する

解法

union-find

このスライドにある解法をそのまま実装した。
http://acm-icpc.aitea.net/index.php?plugin=attach&refer=2009%2FPractice%2F夏合宿%2F講評&openfile=3f.pdf

入力を隣接リストで受け取った後、クエリ ’M' だけを処理する。マークされた頂点を記録する。このとき、同じ部分にマークされた場合、処理せずなかったものとして扱う。

その後、与えられた木を構築するためdfsを行う。このとき、depthに木の深さを記録した。深さの情報をもとに、マークされなかった頂点と、その頂点の親をuniteした。常に根に近い部分を親としてuniteしなければいけないため、深さを利用した。

後はクエリを逆順に処理していく。 ’Q' クエリはfindし、値を足す。 'M' クエリは、自分の親とuniteする。

木を構築する、各頂点の親を持つ、uniteするときに根に近い頂点を親にする、などやることが多く実装に時間がかかった。

コード

#include<bits/stdc++.h>
#define range(i,a,b) for(int i = (a); i < (b); i++)
#define rep(i,b) for(int i = 0; i < (b); i++)
#define all(a) (a).begin(), (a).end()
#define show(x)  cerr << #x << " = " << (x) << endl;
const int INF = 2000000000;
using namespace std;

const int gmax_n = 100005;

int par[gmax_n]; //親
int depth[gmax_n];//木の深さ

void init(int n){
    rep(i,n){
        par[i] = i;
        depth[i] = 0;
    }
}

int find(int x){
    if(par[x] == x){
        return x;
    }else {
        return par[x] = find(par[x]);
    }
}

void unite(int x, int y){
    x = find(x);
    y = find(y);
    if(x == y) return;

    if(depth[x] > depth[y]){
        par[x] = y;
    }else{
        par[y] = x;
    }
}

bool same(int x, int y){
    return find(x) == find(y);
}

void dfs(vector<int> t[gmax_n], int g[gmax_n], int cur){
    rep(i,t[cur].size()){
        int next = t[cur][i];
        if(g[next] != -1) continue;

        depth[next] = depth[cur] + 1;
        g[next] = cur;
        dfs(t,g,next);
    }
}

int main(){
    int n, m;
    while(cin >> n >> m, n||m){
        vector<int> tree[gmax_n];
        init(n);
        rep(i,n - 1){ //木の構築
            int a;
            cin >> a;
            a--;
            tree[i + 1].emplace_back(a);
            tree[a].emplace_back(i + 1);
        }

        bool marked[gmax_n] = {false};//マークされるノードを記録
        marked[0] = true;
        vector<pair<char, int>> q;
        rep(i,m){
            char a;
            int b;
            cin >> a >> b;
            b--;
            if(a == 'M'){
                if(marked[b]) continue; //同じ部分にマークした場合は、初回以外を無視。
                else{
                    marked[b] = true;
                    q.emplace_back(make_pair(a, b));
                }
            }else{
                q.emplace_back(make_pair(a,b));
            }
        }

        int g[gmax_n] = {0};
        range(i,1,gmax_n) g[i] = -1;
        dfs(tree, g, 0);

        rep(i,n){
            if(not marked[i]) unite(g[i], i);
        }

        //rep(i,n) show(g[i])
        //rep(i,n){ show(find(i)) }
        //rep(i,n){ show(depth[i]) }

        long long ans = 0;
        for(int i = q.size() - 1; i >= 0; i--){ //クエリを逆順に処理
            if(q[i].first == 'Q') ans += find(q[i].second) + 1;
            else{
                unite(g[q[i].second], q[i].second);
            }
        }
        cout << ans << endl;
    }
}