概要

\(N\) 個の交差点からなる木構造の町で、毎日新しいお店がオープンします。各日の終わりに、「距離 \(R\) 以内にあるお店がちょうど \(1\) 軒であるような交差点」がいくつあるかを求める問題です。

考察

1. \(R\) の制約に注目する

この問題で最も重要なヒントは、距離の制限である \(R\) が最大 20 と非常に小さい ことです。 もし \(R\) が大きければ、木上のパスや距離に関する高度なデータ構造（重心分解や Heavy-Light Decomposition など）が必要になりますが、\(R \le 20\) であれば、各頂点から距離 \(R\) 以内の情報を直接伝播させるアプローチが有効です。

2. 「嬉しい」条件を整理する

交差点 \(x\) がある日において「嬉しい」と感じるのは、以下の条件を満たすときです。 - 距離 \(R\) 以内にあるお店の中で、最も早くオープンしたお店がすでに開店している。 - 距離 \(R\) 以内にあるお店の中で、2番目に早くオープンしたお店がまだ開店していない。

各お店にはオープンする日（\(1\) 日目〜\(Q\) 日目）が割り当てられています。交差点 \(x\) から距離 \(R\) 以内にあるお店のオープン日のうち、最小値を \(B1_x\)、2番目に小さい値を \(B2_x\) とすると、交差点 \(x\) が嬉しいと感じる期間は \(B1_x\) 日目から \(B2_x - 1\) 日目まで となります。 （お店が1軒もない、あるいは1軒しか存在しない場合は \(B2_x = \infty\) と考えます）

3. \(B1_x, B2_x\) をどう求めるか

動的計画法（DP）のような考え方で、距離を \(1\) ずつ増やしながら情報を広めていくことができます。 - \(dist=0\) のとき：各交差点 \(u\) について、その場所にオープンするお店の番号（日）を \(B1_u\) に格納する。 - \(dist=k\) のとき：距離 \(k-1\) の結果を利用して、距離 \(k\) の結果を更新する。具体的には、頂点 \(u\) の \(B1, B2\) を、隣接する全頂点 \(v\) の距離 \(k-1\) 時点での \(B1, B2\) を使って更新します。

これを \(R\) 回繰り返すことで、各頂点から距離 \(R\) 以内にあるお店の「オープン日の最小値」と「2番目の最小値」を効率的に求めることができます。

アルゴリズム

1. 初期化:

   • 各交差点 \(u\) に対して、その交差点にお店がオープンする日 \(i\) を記録します。お店がない場合は無限大 (\(\infty\)) とします。
   • これを「距離 \(0\) での最小値 \(B1_u\)」とします。

2. 情報の伝播 (DP):

   • 距離 \(d = 1\) から \(R\) まで以下を繰り返します。
   • 各頂点 \(u\) について、自分自身および隣接する頂点 \(v\) が持つ「距離 \(d-1\) での最小値・第2最小値」を集め、その中から改めて最小の2つを選び、頂点 \(u\) の「距離 \(d\) での \(B1_u, B2_u\)」を更新します。

3. 集計 (いもす法):

   • 各交差点 \(u\) について、嬉しい期間 \([B1_u, B2_u - 1]\) がわかります。
   • 全ての交差点についてこの期間を記録するために、長さ \(Q+2\) の差分配列（いもす法）を用意します。
   • diff[B1_u]++, diff[B2_u]-- と操作します。

4. 出力:

   • 差分配列の累積和をとることで、各日 \(i\) における「嬉しい交差点」の総数を出力します。

計算量

• 時間計算量: \(O(R \times (N + M) + Q)\)
  • 情報の伝播に \(R\) 回、全エッジ（\(M = N-1\)）を走査するため \(O(RN)\)。
  • 最後の集計に \(O(N + Q)\)。
  • \(R \le 20\) なので、十分高速に動作します。
• 空間計算量: \(O(N + Q)\)
  • 隣接リスト、および各頂点の \(B1, B2\) を保持する配列が必要です。

実装のポイント

• 2つの最小値の更新: update(b1, b2, x) という関数を作り、現在の最小値 b1 よりも x が小さければ b2 = b1, b1 = x とスライドさせ、b1 より大きく b2 より小さければ b2 = x と更新するようにします。

• 値の入れ替え: 各ステップ \(d\) で一斉に更新するため、現在の値を参照する配列と、次のステップの値を書き込む配列を分けて管理（または swap）する必要があります。

  ソースコード

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

// INF is set to a value larger than any possible shop index (Q <= 2e5)
const int INF = 1e9;

/**
 * Updates the smallest (b1) and second smallest (b2) values with a new value (x).
 * Ensures that b1 and b2 remain distinct and represent the minimums of the set.
 */
inline void update(int& b1, int& b2, int x) {
    if (x >= b2) return;
    if (x < b1) {
        b2 = b1;
        b1 = x;
    } else if (x > b1) {
        b2 = x;
    }
}

int main() {
    // Fast I/O
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q, R;
    if (!(cin >> N >> Q >> R)) return 0;

    // Build the tree as an adjacency list
    vector<vector<int>> adj(N + 1);
    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    // B1[u] and B2[u] store the smallest and second smallest shop indices within distance d from node u.
    // Initially, d = 0, so B1[u] is the index of the shop at u if it exists.
    vector<int> B1(N + 1, INF), B2(N + 1, INF);
    for (int i = 1; i <= Q; ++i) {
        int c;
        cin >> c;
        B1[c] = i;
    }

    // Dynamic Programming to calculate B1 and B2 for distance up to R.
    // In each step d, we update the minimums for distance d using values from distance d-1.
    vector<int> nextB1(N + 1), nextB2(N + 1);
    for (int d = 0; d < R; ++d) {
        for (int u = 1; u <= N; ++u) {
            nextB1[u] = B1[u];
            nextB2[u] = B2[u];
        }
        for (int u = 1; u <= N; ++u) {
            for (int v : adj[u]) {
                // Neighbors' shops at distance d-1 are at distance d from current node u.
                update(nextB1[u], nextB2[u], B1[v]);
                update(nextB1[u], nextB2[u], B2[v]);
            }
        }
        B1.swap(nextB1);
        B2.swap(nextB2);
    }

    // After R steps, B1[u] and B2[u] represent the two smallest indices of shops within distance R of u.
    // Node u is "happy" on day i if exactly one shop within distance R is open.
    // This happens if the first shop (index B1[u]) is open (B1[u] <= i) 
    // and the second shop (index B2[u]) is NOT yet open (B2[u] > i).
    
    // We use a difference array to count happy nodes for each day efficiently.
    vector<int> diff(Q + 2, 0);
    for (int u = 1; u <= N; ++u) {
        if (B1[u] <= Q) {
            diff[B1[u]]++;
            if (B2[u] <= Q) {
                diff[B2[u]]--;
            }
        }
    }

    // Prefix sum of the difference array gives the answer for each day.
    int current_happy_count = 0;
    for (int i = 1; i <= Q; ++i) {
        current_happy_count += diff[i];
        cout << current_happy_count << "\n";
    }

    return 0;
}

この解説は gemini-3-flash-thinking によって生成されました。