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

ゆらのふなびと

競プロ, Python, C++

ABC014 D - 閉路

どこからLCAなんて発想が出てくるの!?と思った人用。

問題

木に辺を追加したときの閉路長を各クエリに対して求めよ。

制約

頂点数:  1 \leq N \leq 10^5

クエリ数:  1 \leq Q \leq 10^5

解法

 (a, b) を追加したときの閉路長は、もとのグラフにおける点  a と点  b との最短距離を  d_{ab} とすると、 d_{ab} + 1 (追加する辺の長さも考えて「+1」)である。 d_{ab} を求めるには「最短経路」という言葉からダイクストラ、ワーシャルフロイドなどの方法が思い浮かぶが、前者はクエリごとに処理を行うので O(Q N \log N ) *1 , 後者は前処理だけで済むが O(N^3)となり間に合わない。そこで、グラフが木であることを利用して最短経路の計算を高速化する。

木においては、ある頂点を「根」と定めることができる。するとDFSにより、各頂点の「深さ」(根からの距離)を求めることができる。頂点vの深さを depth(a)と書くことにしよう。これを使って最短経路を求められないだろうか? 単純に depth(a) + depth(b)とすると、 a, b が根に向かう途中で同じ頂点に合流する場合、そこから先の距離を2回数えてしまうことになる。もし a, b根に向かう途中で初めてぶつかる頂点 cがわかれば、 d_{ab} = depth(a) + depth(b) - 2 * depth(c) と求めることができる。実はこの頂点  c こそLCAである。

ここまでわかれば、あとは公式の解答にあるように、ダブリングを使ってLCAを高速計算すればよい。ダブリングについてはD: 阿弥陀 - AtCoder Beginner Contest 013 | AtCoderが参考になるだろう。LCAについては蟻本p.292にコードがある。

いくつか注意点。(自分がコードを書いていてひっかかったところ)

  • 1番目の親だけはDFSで自分で求める

  • 2, 4, 8,... 番目の親を求めるとき、途中で-1になってしまうケースは場合分け

  • 「ぶつかる直前までa, bを上げる」はステップの大きい方から考えていく(ループを逆順で回す。一種の二分探索)

  • 解答例: C++
#include <bits/stdc++.h>
using namespace std;

#define rep(i,n) for (int i=0;i<(n);i++)
#define rep2(i,a,b) for (int i=(a);i<(b);i++)
#define rrep(i,n) for (int i=(n)-1;i>=0;i--)
#define rrep2(i,a,b) for (int i=(b)-1;i>=(a);i--)
#define pb push_back
#define all(a) (a).begin(),(a).end()

typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<ll> vll;

const int INF = 99999999;
const int MAX_N = 100010;
const int MAX_LOG = 20;

int N, Q;

// par[i][j] = jの2^i番目の親(根を通り過ぎる場合は-1)
int par[MAX_LOG][MAX_N];
int depth[MAX_N];
vector<vi> G(MAX_N);

// 各頂点の深さと、1つ前の親を求める
void dfs(int v, int p, int d) {
    par[0][v] = p;
    depth[v] = d;
    for (auto nxt : G[v]) {
        if (nxt != p) {
            dfs(nxt, v, d + 1);
        }
    }
}

// 各頂点の2^k番目の親を求めておく
void setPar() {
    // 0を根として1つ目の親を求める
    dfs(0, -1, 0);

    // 2^i番目の親を求める
    rep(i, MAX_LOG - 1) {
        rep(j, N) {
            if (par[i][j] == -1) {
                par[i + 1][j] = -1;
            } else {
                par[i + 1][j] = par[i][par[i][j]];
            }
        }
    }
}

int lca(int a, int b) {
    // まずaとbの深さを揃える
    if (depth[a] > depth[b]) {
        swap(a, b);
    }
    rep(i, MAX_LOG) {
        if ((depth[b] - depth[a]) >> i & 1) {
            b = par[i][b];
        }
    }
    if (a == b) return a;

    // ぶつかる直前までa, bを上げる
    rrep(i, MAX_LOG) {
        if (par[i][a] != par[i][b]) {
            a = par[i][a];
            b = par[i][b];
        }
    }
    // aとbの1つ前の親は一致している
    return par[0][a];
}

signed main() {
    cin >> N;
    rep(i, N - 1) {
        int x, y;
        scanf("%d %d", &x, &y);
        x--; y--;
        G[x].emplace_back(y);
        G[y].emplace_back(x);
    }

    setPar();

    cin >> Q;
    rep(i, Q) {
        int a, b;
        cin >> a >> b;
        a--; b--;
        int c = lca(a, b);
        cout << depth[a] + depth[b] - 2 * depth[c] + 1 << endl;
    }
}

ABCのD問題、勉強になるなあ…

*1:ダイクストラの計算量は頂点数を  V , 辺の数を  E とすると  O(V \log E) だが、この問題ではグラフは木構造、すなわち  V = N, E = N - 1 なのでダイクストラの計算量は  O(N \log N) 。これをクエリごとに行うので  O(Q N \log N )