実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 816

問題概要

$N$ 個の頂点からなる木が与えられ、 $i$ 番目の辺は頂点 $A_i$ と $B_i$ を結んでいる。この木の部分木のうち、指定された $K$ 個の頂点 $V_1, V_2, \cdots, V_K$ をすべて含むようなものの頂点数の最小値を求めよ。

制約

入力はすべて整数。
$1 \leq K \leq N \leq 2 \times 10^5$
$1 \leq A_i, B_i \leq N$
$1 \leq V_1 \leq V_2 \leq \cdots \leq V_K \leq N$

考察

タイトルにある通り、求めるべき木は「最小シュタイナー木」と呼ばれるものである。ただ、本来の最小シュタイナー木は辺に重みが付いたものであり、インターネットに転がっているアルゴリズムは $K = 11$ 程度が限界らしく、本問の制約では間に合わない。

では、ここからは「残すべき頂点とはどのような頂点か？」を考えてみよう。

与えられた木 $T$ の中の任意の頂点 $v$ について、 $v$ を根とする部分木を $T_v$ とする。実は、 $T_v$ の中の任意の頂点 $u$ が $V_1, V_2, \cdots, V_K$ のいずれかに該当する場合、頂点 $v$ は削除できない頂点である。

なぜなら、 $u, v$ 間には辺がただ一つしか存在しないので、 $v$ を削除すると $u$ に到達できなくなってしまうからである。

逆に、そうでない頂点 $v$ は削除してよいことになる。

これをどう実装していくかだが、簡単なのは $V_1$ から DFS で頂点を走査していく方法だろう。

コードではcntという配列を用意して、cnt[i]に頂点 $i$ を根とする部分木の頂点の中で $V_1, V_2, \cdots, V_K$ に含まれているものの個数を格納している。

$\mathrm{cnt}_i \geq 1$ となる頂点 $i$ は残すべき頂点なので、最後にこの個数を出力すればよい。

コード

#include <bits/stdc++.h>
using namespace std;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

template <class T>
using Graph = vector<vector<T>>;

// ======================================== //

void dfs(Graph<int> &G, set<int> &V, vector<int> &cnt, int now, int pre)
{
    if (V.count(now)) cnt[now]++;

    for (auto &v : G[now]) {
        if (v == pre) continue;

        dfs(G, V, cnt, v, now);
        cnt[now] += cnt[v];
    }
}

int main()
{
    int N, K;
    cin >> N >> K;
    Graph<int> G(N);
    rep(i, 0, N - 1) {
        int A, B;
        cin >> A >> B;
        G[A - 1].push_back(B - 1);
        G[B - 1].push_back(A - 1);
    }
    set<int> V;
    rep(i, 0, K) {
        int v;
        cin >> v;
        V.insert(v - 1);
    }

    vector<int> cnt(N, 0);
    dfs(G, V, cnt, *V.begin(), -1);

    int ans = 0;
    for (auto &x : cnt) {
        if (x >= 1) ans++;
    }

    cout << ans << endl;
}