Official

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

gemini-3-flash-thinking

概要

この問題は、木構造(組織図)において特定の頂点(部署) \(K\) を根とする部分木のサイズを求め、全体の頂点数 \(N\) から差し引く問題です。

考察

問題の整理

  • 会社全体は \(N\) 個の部署からなる木構造です。
  • 部署 \(1\) が根(本社)です。
  • 部署 \(K\) を解体すると、部署 \(K\) 自身と、その配下にあるすべての部署(子、孫、……)が消滅します。
  • 求めたいのは「残った部署の数」、つまり \(N - (\text{部署 } K \text{ を根とする部分木のサイズ})\) です。

部分木のサイズをどう求めるか

通常、木構造において部分木のサイズを求めるには、深さ優先探索(DFS)などを用いて葉の方から順にカウントしていく手法が一般的です。しかし、この問題には非常に重要な制約があります。

\(1 \leq P_i < i\)\(2 \leq i \leq N\)

この制約は、「親の番号は必ず子の番号よりも小さい」 ということを意味しています。 つまり、番号が大きい部署から順に処理していけば、ある部署 \(i\) を処理する時点で、その子となる部署(番号が \(i\) より大きいもの)のサイズはすでに計算し終わっていることが保証されます。

この性質を利用すると、複雑なグラフ探索(DFSやBFS)を実装することなく、単純なループだけで部分木のサイズを計算することができます。

アルゴリズム

動的計画法(DP)に近い考え方で、以下の手順で解きます。

  1. 初期化: 各部署 \(i\) の部分木サイズを保持する配列 subtree_size を用意し、すべて \(1\) で初期化します(自分自身の分)。
  2. ボトムアップの集計: 部署番号 \(i\)\(N\) から \(2\) まで逆順にループを回します。
    • 部署 \(i\) の親を \(P_i\) とすると、subtree_size[P_i]subtree_size[i] を加算します。
    • 番号の大きい順に処理することで、親に値を足し合わせる前に、その子のサイズが確定している状態を作れます。
  3. 答えの計算: 全体の数 \(N\) から、解体対象である部署 \(K\) の部分木サイズ subtree_size[K] を引いた値を出力します。

計算量

  • 時間計算量: \(O(N)\) 部署 \(N\) から \(2\) までを \(1\) 回ずつ走査するため、頂点数 \(N\) に対して線形時間で計算が終わります。
  • 空間計算量: \(O(N)\) 親の情報を格納する配列や、部分木のサイズを格納する配列に \(O(N)\) のメモリを使用します。

実装のポイント

  • 入力の高速化: \(N\) が最大 \(2 \times 10^5\) と大きいため、Pythonでは sys.stdin.read().split() を使うなどして、大量の入力を一括で読み込むのが効率的です。

  • インデックスの調整: 問題文の \(P_i\) は部署 \(2, 3, \ldots, N\) の親の情報が順に与えられます。リスト P_list に格納した場合、部署 \(i\) の親は P_list[i-2] でアクセスすることになる点に注意しましょう。

  • 逆順ループ: range(N, 1, -1) を使うことで、部署 \(N\) から順に親へとサイズを伝播させることができます。これにより、再帰関数の呼び出しによるスタックオーバーフローの心配もありません。

    ソースコード

import sys

def solve():
    # 競技プログラミングにおいて、大量の入力を高速に読み込むために sys.stdin.read を使用します
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # N: 部署の総数, K: 解体する部署の番号
    N = int(input_data[0])
    K = int(input_data[1])
    
    # P_list は番号 2 から N までの各部署の直属の上位部署のリストです
    # P_list[0] が P_2, P_list[N-2] が P_N に対応します
    if N > 1:
        P_list = list(map(int, input_data[2:]))
    else:
        P_list = []
        
    # subtree_size[i] は部署 i を根とする部分木に含まれる部署の数を保持します
    # 最初は自分自身のみをカウントするため、すべて 1 で初期化します
    subtree_size = [1] * (N + 1)
    
    # 制約 P_i < i により、番号が大きい部署から順に処理することで、
    # ある部署を処理する際には、その部署のすべての子(下部組織)のサイズが計算済みであることが保証されます
    # 番号 N から 2 まで逆順にループを回します
    for i in range(N, 1, -1):
        # 部署 i の親(上位部署)の番号を取得
        # P_list は 0-indexed なので、部署 i の親は P_list[i-2] に格納されています
        parent = P_list[i-2]
        # 親の subtree_size に、現在の部署 i の subtree_size を加算します
        subtree_size[parent] += subtree_size[i]
        
    # 全体の部署数 N から、解体対象の部署 K とその傘下の部署の総数(subtree_size[K])を引きます
    ans = N - subtree_size[K]
    
    # 結果を出力します
    print(ans)

if __name__ == '__main__':
    solve()

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

posted:
last update: