概要

各交差点について、「距離 \(R\) 以下にあるお店」のうち、最も早く開店するお店と 2 番目に早く開店するお店を求めます。
その情報が分かれば、その交差点が「嬉しい」期間を区間として表せるので、差分配列で各日の答えを求められます。

考察

ある交差点 \(x\) に注目します。

距離 \(R\) 以下にあるお店の開店日を小さい順に

\[ t_1, t_2, t_3, \dots \]

とします。

このとき、\(i\) 日目の終わりに距離 \(R\) 以下で開いているお店の数は、\(t_j \leq i\) であるものの個数です。

したがって、交差点 \(x\) が嬉しいのは、

• 最初のお店が開店した日 \(t_1\) 以降
• 2 番目のお店が開店する日 \(t_2\) より前

です。

つまり、嬉しい日は

\[ t_1 \leq i < t_2 \]

です。

もし距離 \(R\) 以下にお店が 1 軒しかないなら、\(t_1\) 日目から最終日 \(Q\) 日目までずっと嬉しいです。

よって、各交差点について必要なのは、

• 距離 \(R\) 以下にあるお店の最小開店日
• 距離 \(R\) 以下にあるお店の 2 番目に小さい開店日

の 2 つだけです。

例えば、ある交差点から距離 \(R\) 以下のお店が 3 日目、7 日目、10 日目に開店するとします。
この交差点が嬉しいのは 3 日目から 6 日目までです。

これは差分配列で

• diff[3] += 1
• diff[7] -= 1

とすれば表せます。

素朴に各日ごとに新しく開店したお店から距離 \(R\) 以下の交差点をすべて調べると、木が星型のような場合に 1 回の探索で \(O(N)\) かかることがあります。
これを \(Q\) 日分行うと \(O(NQ)\) となり、制約では間に合いません。

そこで、すべての開店日を先に見ておき、各交差点について「距離 \(R\) 以下の開店日の最小 2 つ」をまとめて求めます。

ここで重要なのは、\(R \leq 20\) と小さいことです。
辺を通じて「最小 2 つの開店日」を \(R\) 回伝播させれば、距離 \(R\) 以下の情報を得られます。

アルゴリズム

まず、各交差点 \(v\) について次の 2 つを管理します。

• best1[v]: 現在分かっている範囲で、最も早い開店日
• best2[v]: 現在分かっている範囲で、2 番目に早い開店日

最初は距離 \(0\)、つまり自分自身だけを見ます。

• 交差点 \(C_i\) にお店が開くなら、best1[C_i] = i
• お店が開かない交差点は INF

とします。

次に、これを \(R\) 回繰り返します。

1 回の操作で、各辺 \((u, v)\) について、

• \(v\) が知っている最小 2 つの開店日を \(u\) に伝える
• \(u\) が知っている最小 2 つの開店日を \(v\) に伝える

という処理を行います。

これにより、1 回目の後には距離 \(1\) 以下、2 回目の後には距離 \(2\) 以下、というように情報が広がります。
\(R\) 回繰り返すと、各交差点について距離 \(R\) 以下にあるお店の開店日の最小 2 つが求まります。

候補を追加するときは、同じ開店日を重複して数えないようにします。
同じお店の情報が複数経路で伝わる可能性があるためです。
この問題では 1 日に 1 軒だけお店が開くので、開店日はお店を一意に表します。

最後に、各交差点 \(x\) について、

• first = best1[x]
• second = best2[x]

とします。

first が存在しない場合、その交差点の距離 \(R\) 以下にはお店がないので、答えに影響しません。

first が存在する場合、

• diff[first] += 1
• second が存在するなら diff[second] -= 1

とします。

その後、差分配列の累積和を取ることで、各日の答えが求まります。

計算量

• 時間計算量: \(O(RN + Q)\)
• 空間計算量: \(O(N + Q)\)

\(R \leq 20\) なので、\(O(RN)\) は十分高速です。

実装のポイント

best1[v], best2[v] を更新するときは、「最小の異なる 2 つの値」を保つ必要があります。

例えば、新しい候補 t を追加する関数では、次のようにします。

• t が INF なら無視
• t がすでに best1 または best2 に入っているなら無視
• t < best1 なら、best2 = best1, best1 = t
• そうでなく t < best2 なら、best2 = t

また、伝播処理では同じステップ内の更新がさらに連鎖しないように、現在の配列から次の配列 nxt を作ってから入れ替えます。
これにより、1 回のループでちょうど距離が 1 だけ広がることが保証されます。

ソースコード

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

const int INF = 1e9;

inline void add_candidate(int &a, int &b, int t) {
    if (t == INF) return;
    if (t == a || t == b) return;

    if (t < a) {
        b = a;
        a = t;
    } else if (t < b) {
        b = t;
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q, R;
    cin >> N >> Q >> R;

    vector<pair<int, int>> edges;
    edges.reserve(max(0, N - 1));

    for (int i = 0; i < N - 1; i++) {
        int u, v;
        cin >> u >> v;
        --u; --v;
        edges.emplace_back(u, v);
    }

    vector<int> best1(N, INF), best2(N, INF);

    for (int i = 1; i <= Q; i++) {
        int c;
        cin >> c;
        --c;
        best1[c] = i;
    }

    for (int step = 0; step < R; step++) {
        vector<int> nxt1 = best1;
        vector<int> nxt2 = best2;

        for (auto [u, v] : edges) {
            add_candidate(nxt1[u], nxt2[u], best1[v]);
            add_candidate(nxt1[u], nxt2[u], best2[v]);

            add_candidate(nxt1[v], nxt2[v], best1[u]);
            add_candidate(nxt1[v], nxt2[v], best2[u]);
        }

        best1.swap(nxt1);
        best2.swap(nxt2);
    }

    vector<long long> diff(Q + 3, 0);

    for (int x = 0; x < N; x++) {
        int first = best1[x];
        int second = best2[x];

        if (first <= Q) {
            diff[first]++;
            if (second <= Q) diff[second]--;
        }
    }

    long long ans = 0;
    for (int i = 1; i <= Q; i++) {
        ans += diff[i];
        cout << ans << '\n';
    }

    return 0;
}

---

この解説は gpt-5.5-xhigh によって生成されました。