Official

C - 部署の解体 / Dissolution of the Department Editorial by admin

GPT 5.2 High

概要

木で表された部署構造から、指定された部署 \(K\) を根とする部分木(\(K\) とその配下すべて)を削除したときに残る部署数を求めます。

考察

この問題で削除されるのは「部署 \(K\) の部分木に含まれる頂点すべて」です。したがって、

  • 残る部署数 = 全体の部署数 \(N\) − (\(K\) の部分木サイズ)

に帰着します。

重要なのは「\(K\) の部分木サイズ」を正しく数えることです。入力は各頂点 \(i(2 \le i \le N)\) について親 \(P_i\) が与えられる形なので、まず親から子への辺(子リスト)を作っておくと、\(K\) から下へ辿る探索(DFS/BFS)で部分木サイズを数えられます。

素朴に「各頂点が \(K\) の配下かどうか」を親を辿って判定すると、最悪で各頂点ごとに根まで遡って \(O(N)\)、全体で \(O(N^2)\) になり \(N \le 2 \times 10^5\) では間に合いません。
そこで、子リストを作って \(K\) から一度だけ探索し、部分木内の頂点だけを \(O(\text{部分木サイズ})\) で数えます。

例:
\(K=3\) の部分木に 5 頂点あるなら、削除後は \(N-5\) 個が残ります。

アルゴリズム

  1. 入力の親配列 \(P_2,\dots,P_N\) を使い、各頂点の子を列挙する配列 children を作る。
    • children[p] に子 i を追加する(辺 \(p \to i\) を張るイメージ)。
  2. スタック(またはキュー)を用意し、\(K\) を入れて探索する(反復DFS)。
  3. スタックから頂点 \(v\) を取り出すたびにカウントを 1 増やし、children[v] にある子をスタックに追加する。
  4. 探索で数えた個数 sub が部分木サイズなので、答えは \(N - sub\)

計算量

  • 時間計算量: \(O(N)\)
    (子リスト作成が \(O(N)\)、探索も各頂点を高々 1 回処理)
  • 空間計算量: \(O(N)\)
    (子リストとスタック分)

実装のポイント

  • 再帰DFSは深さが最大 \(N\) になりうるため、Python では再帰上限に引っかかる可能性があります。コードではスタックを使う反復DFSで安全に実装しています。

  • P\(P_2\) から始まるので、頂点 i の親は P[i-2] になる点に注意します。

  • \(N=1\) のときは親情報の行が空になるため、P を空配列として扱えるようにしています。

    ソースコード

import sys

def main():
    input = sys.stdin.buffer.readline
    N, K = map(int, input().split())
    P = list(map(int, input().split())) if N >= 2 else []

    children = [[] for _ in range(N + 1)]
    for i in range(2, N + 1):
        p = P[i - 2]
        children[p].append(i)

    stack = [K]
    sub = 0
    while stack:
        v = stack.pop()
        sub += 1
        stack.extend(children[v])

    print(N - sub)

if __name__ == "__main__":
    main()

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

posted:
last update: