Official

L - 消防署と緊急出動 / Fire Station and Emergency Dispatch Editorial by admin

gpt-5.5-high

概要

開設中の消防署が動的に増減する重み付き木上で、ある町から最も近い開設中の消防署までの距離が \(T\) 以下かを判定する問題です。

各問い合わせを高速に処理するため、木を重心分解し、各重心ごとに「その重心から見た最寄りの開設中消防署までの距離」を管理します。

考察

求めたいものは、各タイプ 2 V T に対して

\[ \min_{x \in \text{開設中の町}} \mathrm{dist}(V, x) \]

が \(T\) 以下かどうかです。

素朴な方法では間に合わない

例えば、各問い合わせごとに町 \(V\) から DFS / Dijkstra をして最寄りの開設中消防署を探すと、\(1\) 回あたり \(O(N)\) かかります。

\(N, Q \leq 5 \times 10^4\) なので、最悪で

\[ O(NQ) \]

となり、間に合いません。

また、開設中の消防署をすべて走査して距離を計算する方法も、開設中の数が多いと同様に間に合いません。

重心分解を使う

木に対する「動的な最寄り頂点クエリ」では、重心分解が有効です。

重心分解では、木の重心を取り除き、残った部分木に対して再帰的に同じことを行います。

このとき、各頂点 \(v\) について、重心分解上で関係する重心たち \(c\) と、その距離

\[ \mathrm{dist}(v, c) \]

をあらかじめ保存しておきます。

重心分解では、各頂点が関わる重心の数は \(O(\log N)\) 個です。

距離の分解

任意の \(2\) 頂点 \(u, v\) について、重心分解上で両者を分ける重心 \(c\) が存在します。

そのような \(c\) については、元の木上の \(u\) から \(v\) へのパスが \(c\) を通るため、

\[ \mathrm{dist}(u, v) = \mathrm{dist}(u, c) + \mathrm{dist}(v, c) \]

となります。

したがって、町 \(v\) から最も近い開設中消防署までの距離は、

\[ \min_{c \in \text{paths}[v]} \left( \mathrm{dist}(v, c) + \min_{x \in \text{開設中},\ c \in \text{paths}[x]} \mathrm{dist}(x, c) \right) \]

で求められます。

つまり、各重心 \(c\) ごとに

現在開設中の消防署 \(x\) についての \(\mathrm{dist}(x, c)\) の最小値

を管理できればよいです。

アルゴリズム

1. 重心分解を行う

各町 \(v\) について、以下の情報を保存します。

\[ (c, \mathrm{dist}(v, c)) \]

ここで \(c\) は、重心分解において \(v\) が属する部分木の重心です。

コードでは paths[v] にこの情報を保存しています。

例えば、

paths[v] = [c1, d1, c2, d2, ...]

のように、重心と距離を交互に格納しています。

各頂点について保存される重心の数は \(O(\log N)\) 個です。

2. 各重心ごとに開設中消防署までの距離を管理する

各重心 \(c\) について、以下の多重集合を管理したいです。

\[ \{\mathrm{dist}(x, c) \mid x \text{ は開設中の消防署}\} \]

この最小値が分かれば、問い合わせ時に使えます。

ただし、消防署は閉鎖されることもあるため、値の削除が必要です。

Python の heapq は最小値の取得は高速ですが、任意の値の削除は苦手です。

そこで、以下の \(2\) つのヒープを使います。

  • add_heaps[c]:追加された距離
  • del_heaps[c]:削除予定の距離

値を削除したいときは、実際には del_heaps[c] に入れておきます。

最小値を見る直前に、

while add_heaps[c] and del_heaps[c] and add_heaps[c][0] == del_heaps[c][0]:
    heappop(add_heaps[c])
    heappop(del_heaps[c])

として、削除済みの値をまとめて取り除きます。

これは遅延削除と呼ばれるテクニックです。

3. 初期状態

初期状態では町 \(1\) の消防署のみが開設中です。

そのため、paths[1] に含まれるすべての重心 \(c\) に対して、

\[ \mathrm{dist}(1, c) \]

を add_heaps[c] に追加します。

4. タイプ 1 X

町 \(X\) の消防署の状態を切り替えます。

閉鎖中なら開設

paths[X] に含まれる各重心 \(c\) について、

\[ \mathrm{dist}(X, c) \]

を add_heaps[c] に追加します。

開設中なら閉鎖

同様に、各重心 \(c\) について、

\[ \mathrm{dist}(X, c) \]

を del_heaps[c] に追加します。

実際の削除は、後でその重心の最小値を見るときに行います。

5. タイプ 2 V T

町 \(V\) から最も近い開設中消防署までの距離が \(T\) 以下かを判定します。

paths[V] に含まれる各重心 \(c\) について、次を考えます。

\[ \mathrm{dist}(V, c) + \min_{x \text{ 開設中}} \mathrm{dist}(x, c) \]

もしこれが \(T\) 以下なら、答えは YES です。

コードでは、

rem = T - dist(V, c)

として、重心 \(c\) から見て距離 rem 以下の開設中消防署が存在するかを確認しています。

すべての重心について条件を満たさなければ NO です。

計算量

  • 時間計算量: \(O(N \log N + Q \log N \log (N+Q))\) 償却
  • 空間計算量: \(O((N+Q) \log N)\)

重心分解により、各頂点が持つ重心の個数は \(O(\log N)\) 個です。

各イベントでは高々 \(O(\log N)\) 個の重心だけを見ればよく、ヒープ操作に \(O(\log (N+Q))\) かかります。

また、遅延削除で取り除かれる値は、それぞれ高々一度だけ pop されるため、全体として償却計算量で抑えられます。

実装のポイント

  • 距離の最大値は大きくなるため、C++ などでは long long が必要です。

    • Python では整数の桁あふれを気にする必要はありません。
  • 町 \(1\) の消防署は常に開設中で、入力でもタイプ 1 に町 \(1\) は現れません。

  • heapq では任意削除ができないため、add_heaps と del_heaps による遅延削除を使います。

  • paths[v] は (重心, 距離) の組を保存します。

    • コードでは高速化のため、タプルのリストではなく [c, d, c, d, ...] の形で保存しています。
  • 重心分解では再帰を使うため、Python では sys.setrecursionlimit を大きめに設定しておくと安全です。

    ソースコード

import sys
from heapq import heappush, heappop

def main():
    input = sys.stdin.buffer.readline
    sys.setrecursionlimit(200000)

    N, Q = map(int, input().split())
    g = [[] for _ in range(N + 1)]
    for _ in range(N - 1):
        a, b, c = map(int, input().split())
        g[a].append((b, c))
        g[b].append((a, c))

    blocked = [False] * (N + 1)
    parent = [0] * (N + 1)
    size = [0] * (N + 1)
    paths = [[] for _ in range(N + 1)]

    def decompose(start):
        order = []
        stack = [start]
        parent[start] = 0

        while stack:
            v = stack.pop()
            order.append(v)
            pv = parent[v]
            for to, _ in g[v]:
                if to != pv and not blocked[to]:
                    parent[to] = v
                    stack.append(to)

        for v in reversed(order):
            s = 1
            pv = parent[v]
            for to, _ in g[v]:
                if to != pv and not blocked[to]:
                    s += size[to]
            size[v] = s

        total = len(order)
        half = total // 2
        centroid = -1

        for v in order:
            mx = total - size[v]
            for to, _ in g[v]:
                if not blocked[to] and parent[to] == v:
                    if size[to] > mx:
                        mx = size[to]
            if mx <= half:
                centroid = v
                break

        c = centroid
        blocked[c] = True

        stack = [(c, 0, 0)]
        while stack:
            v, p, d = stack.pop()
            paths[v].append(c)
            paths[v].append(d)
            for to, w in g[v]:
                if to != p and not blocked[to]:
                    stack.append((to, v, d + w))

        for to, _ in g[c]:
            if not blocked[to]:
                decompose(to)

    decompose(1)

    g = blocked = parent = size = None

    add_heaps = [[] for _ in range(N + 1)]
    del_heaps = [[] for _ in range(N + 1)]
    active = [False] * (N + 1)
    active[1] = True

    p = paths[1]
    for i in range(0, len(p), 2):
        heappush(add_heaps[p[i]], p[i + 1])

    out = []

    for _ in range(Q):
        s = input().split()

        if s[0] == b'1':
            x = int(s[1])
            p = paths[x]

            if active[x]:
                active[x] = False
                for i in range(0, len(p), 2):
                    heappush(del_heaps[p[i]], p[i + 1])
            else:
                active[x] = True
                for i in range(0, len(p), 2):
                    heappush(add_heaps[p[i]], p[i + 1])

        else:
            v = int(s[1])
            T = int(s[2])
            p = paths[v]
            ok = False

            i = 0
            lp = len(p)
            while i < lp:
                c = p[i]
                d = p[i + 1]
                rem = T - d

                if rem >= 0:
                    ah = add_heaps[c]
                    dh = del_heaps[c]

                    while ah and dh and ah[0] == dh[0]:
                        heappop(ah)
                        heappop(dh)

                    if ah and ah[0] <= rem:
                        ok = True
                        break

                i += 2

            out.append("YES" if ok else "NO")

    sys.stdout.write("\n".join(out))

if __name__ == "__main__":
    main()

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

posted:
last update: