Aizu Online Judge 2710: Marked Ancestor

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2170

問題概要

N頂点の木とQ個のクエリが与えられる。木は常に頂点1が根である。また初めは頂点1だけにマークがされている。クエリにはMとQの2種類があり、Mクエリでは指定された頂点にマークを付ける。Qクエリでは指定された頂点の祖先でマークが付いているもののうちその頂点から最も近い頂点を求める(祖先には自分も含む)。Qクエリで求めた頂点番号の合計を出力せよ。

N, Q <= 105

解法

マークを付ける操作を考えると、これは「マークを付けた頂点を根とする部分木を木から切り離す」操作ともいえる。切り離された部分木に含まれるノードたちにとっては、その時点での最も近いマークされた祖先はその部分木の根である。

このように見るとこれはもともとひとつだった集まりをクエリごとに分断していき、途中途中でその集まりの親を聞かれている問題と捉えることができる。このような操作を要求される問題に対する典型テクとして「操作を逆順で見ていく」というのがあり、この問題にもそれを使うことができる。

なぜ逆順で見るのがいいかというと、「集合を切って分割していく操作」というのは扱いづらいが、これを逆から見た「分割されている集合たちを繋いでいく操作」はかなり扱いやすい性質を持っているからである。具体的にはUnionFindが使える。

というわけでクエリを逆回しに見ていくことにしたいので、まずクエリを先に読んですべてのマークがついたことにしてしまう。そしてこの時点でマークがついていないノードはすべてそのノードの親とuniteする。このようにしてできる集合たちが上述の「部分木を切り離していく操作」を最後まで行ったときのそれぞれの部分木に対応する。あとはここからまたクエリを逆順に見ていく。

逆順に見ていってQクエリが飛んできた場合、指定された頂点が属する集合の根に当たる頂点、すなわち最も深さが小さい頂点を答えればよい。これはUnionFindを適当に改造してそのような情報を持たせることにすれば実装できる。

Mクエリが飛んできた場合、その頂点からマークを剥がすことになる。これはつまり部分木を元の木にくっつける作業で、マークを剥がす頂点とその直接の親をuniteすればよい。ただしその頂点に何度もマークが行われているときは注意が必要で、この場合いちばん最初につけられたマーク以外には意味がないため、他は無視しないといけない。

感想

AOJ-ICPC軽く埋めたろかと思って見たけど普通にむずくて解けなかった

木の頂点をUFで管理するという発想がなかったのは反省

コード(D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.string;

bool solve() {
    auto s = readln.split.map!(to!int);
    auto N = s[0];
    auto Q = s[1];
    if (N == 0) return false;

    auto G = new int[][](N);
    auto D = new int[](N);
    auto P = new int[](N);
    P[0] = -1;

    foreach (i; 1..N) {
        auto n = readln.chomp.to!int;
        P[i] = n-1;
        G[i] ~= n-1;
        G[n-1] ~= i;
    }

    void dfs(int n, int p, int d) {
        D[n] = d;
        foreach (m; G[n]) if (m != p) dfs(m, n, d+1);
    }

    dfs(0, -1, 0);

    Tuple!(int, int)[] queries;
    auto marked = new int[](N);

    foreach (i; 1..Q+1) {
        auto t = readln.split;
        auto n = t[1].to!int - 1;
        if (t[0] == "M" && !marked[n]) {
            queries ~= tuple(0, n);
            marked[n] = i;
        } else if (t[0] == "Q") {
            queries ~= tuple(1, n);
        }
    }

    auto uf = new UnionFind(N, D);
    foreach (i; 0..N) if (!marked[i]) uf.unite(i, P[i]);
    long ans = 0;

    foreach_reverse (q; queries) {
        if (q[0] == 0) {
            uf.unite(q[1], P[q[1]]);
        } else {
            ans += uf.find(q[1]) + 1;
        }
    }

    ans.writeln;

    return true;
}

void main() {
    while (solve()) {}
}


class UnionFind {
    int N;
    int[] table;
    int[] depth;

    this(int n, int[] d) {
        N = n;
        table = new int[](N);
        fill(table, -1);
        depth = d.dup;
    }

    int find(int x) {
        return table[x] < 0 ? x : (table[x] = find(table[x]));
    }

    void unite(int x, int y) {
        if (x == -1 || y == -1) return;
        x = find(x);
        y = find(y);
        if (x == y) return;
        if (depth[x] > depth[y]) swap(x, y);
        table[x] += table[y];
        table[y] = x;
    }
}