けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 014 D - 閉路 (青色)

木上のパスに関する問題!!
LCA で解決できる典型問題

問題概要

頂点数  N の木が与えられる。次の  Q 個のクエリに答えよ。

  • 各クエリでは木上の 2 頂点  a, b が与えられる
  • 木に辺  (a, b) を仮に追加したとすると、閉路が 1 個形成される
  • その閉路に含まれる辺の本数を答えよ

制約

  •  1 \le N, Q \le 10^{5}

考えたこと

グラフ理論的には、木とは「閉路をもたないグラフ」ということになる。下図のように、それに辺を新たに付け加えると閉路ができるのだ。

一般に、辺  (a, b) を付け加えることでできる閉路の長さは、パス  a- b の長さに 1 を足したものになる。

よって各クエリは「パス  a- b の長さを答えてください (それに 1 を足して出力してください)」と言い換えられるのだ。

f:id:drken1215:20210619154232p:plain

木のパスの長さ

木のパスの長さに答える問題は超有名問題で、蟻本上級編「グラフマスターへの道」にも解説がある。いろんなやり方があるが、LCA を求める方法が有名だ (さらに LCA を求める方法もたくさんある)。

まず、木の根を 1 つ決めて根付き木にしてしまう。そして、二頂点  a, b に対して、LCA (最近共通祖先) を求めよう (下図の頂点  v)。

LCA の求め方は蟻本参照。

f:id:drken1215:20210619174729p:plain

LCA が求められたら、次の値を合計することでパスの長さが求められる。

  • これはパス  a- v の長さ
  • これはパス  b- v の長さ

前者は (頂点  a の深さ) - (頂点  v の深さ) で求められて、後者は (頂点  b の深さ) - (頂点  v の深さ) で求められる。つまり、


depth[a] + depth[b] - 2 * depth[v]


で求められる (depth は各頂点の深さ)。

コード

ここでは、LCA を求める処理をライブラリ化したものを用いた。

github.com

計算量は  O(N + Q \log N) となる。

#include <iostream>
#include <vector>
using namespace std;

// グラフを表す構造体
using Graph = vector<vector<int>>;

// LCA を求めるライブラリ
struct LCA {
    vector<vector<int> > parent; // parent[d][v] := 2^d-th parent of v
    vector<int> depth;
    LCA() { }
    LCA(const Graph &G, int r = 0) { init(G, r); }
    void init(const Graph &G, int r = 0) {
        int V = (int)G.size();
        int h = 1;
        while ((1<<h) < V) ++h;
        parent.assign(h, vector<int>(V, -1));
        depth.assign(V, -1);
        dfs(G, r, -1, 0);
        for (int i = 0; i+1 < (int)parent.size(); ++i)
            for (int v = 0; v < V; ++v)
                if (parent[i][v] != -1)
                    parent[i+1][v] = parent[i][parent[i][v]];
    }
    void dfs(const Graph &G, int v, int p, int d) {
        parent[0][v] = p;
        depth[v] = d;
        for (auto e : G[v]) if (e != p) dfs(G, e, v, d+1);
    }
    int get(int u, int v) {
        if (depth[u] > depth[v]) swap(u, v);
        for (int i = 0; i < (int)parent.size(); ++i)
            if ( (depth[v] - depth[u]) & (1<<i) )
                v = parent[i][v];
        if (u == v) return u;
        for (int i = (int)parent.size()-1; i >= 0; --i) {
            if (parent[i][u] != parent[i][v]) {
                u = parent[i][u];
                v = parent[i][v];
            }
        }
        return parent[0][u];
    }
};

int main() {
    int N;
    cin >> N;
    Graph G(N);
    for (int i = 0; i < N-1; ++i) {
        int a, b;
        cin >> a >> b;
        --a, --b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    // LCA ライブラリ
    LCA lca(G);

    // クエリ処理
    int Q;
    cin >> Q;
    for (int q = 0; q < Q; ++q) {
        int a, b;
        cin >> a >> b;
        --a, --b;

        // lca
        int v = lca.get(a, b);

        // パスの長さ
        int len = lca.depth[a] + lca.depth[b] - 2 * lca.depth[v];
        cout << len + 1 << endl;
    }
}