Official

C - 組織のアンバランス / Imbalance of the Organization Editorial by admin Editorial for printing

Gemini 3.0 Flash (Thinking)

概要

この問題は、根付き木として表される組織図において、各社員の「直属の部下たちが率いるチームの仕事量(部分木の重み合計)」の最大値と最小値の差を求め、その全社員における最大値を答える問題です。

考察

この問題を解くためには、各社員について以下の2つの情報を知る必要があります。 1. チームの仕事量: その社員を根とする部分木に含まれる全社員の仕事量 \(W_i\) の総和。 2. 直属の部下たちのチーム仕事量の最大値と最小値: アンバランス度を計算するために必要。

重要な制約の活用

問題文にある \(P_i \leq i - 1\) という制約が非常に重要です。これは「上司の番号は常に部下の番号より小さい」ことを意味しており、社員の番号順が一種のトポロジカル順序(根から葉への順序)になっていることを示しています。

この性質を利用すると、高価なグラフ探索(DFSやBFS)を行わなくても、社員番号の大きい順(\(N\) から \(2\) まで)に処理を行うだけで、ボトムアップ(葉から根へ)の集計が可能になります。社員 \(i\) を処理する時点で、その部下(番号が \(i\) より大きい社員)の集計はすべて完了していることが保証されるからです。

素朴なアプローチ

もし各社員ごとに毎回部分木を走査して合計を計算すると、最悪の場合 \(O(N^2)\) の時間がかかり、 \(N = 2 \times 10^5\) では間に合いません。そのため、動的計画法(DP)のように、下の階層から計算結果を積み上げていく必要があります。

アルゴリズム

以下の手順で解を求めます。

  1. 初期化:

    • 各社員のチーム仕事量 \(S_i\) を、その社員自身の仕事量 \(W_i\) で初期化します。
    • 各社員の「部下の仕事量の最大値 child_max」「最小値 child_min」「部下の数 child_count」を管理する配列を用意します。
  2. ボトムアップの集計:

    • 社員番号 \(i\) を \(N\) から \(2\) まで逆順にループします。
    • 社員 \(i\) の現在のチーム仕事量 \(S_i\) を、その上司 \(P_i\) のチーム仕事量 \(S_{P_i}\) に加算します。
    • 同時に、上司 \(P_i\) における child_max, child_min, child_count を \(S_i\) を使って更新します。
  3. アンバランス度の計算:

    • 全ての社員について、child_count が 2 以上のもののうち、child_max - child_min の最大値を求めます。

計算量

  • 時間計算量: \(O(N)\)
    • 入力の読み込みに \(O(N)\)、ボトムアップの集計ループで \(O(N)\)、最終的な最大値の探索で \(O(N)\) となり、全体で線形時間で動作します。
  • 空間計算量: \(O(N)\)
    • 社員ごとの仕事量や上司情報を保持するために、長さ \(N\) の配列を数個使用します。

実装のポイント

  • 効率的な入力: Pythonで \(N=2 \times 10^5\) 規模の入力を扱う場合、input() を繰り返すよりも sys.stdin.read().split() で一括で読み込む方が高速です。

  • 1-indexed と 0-indexed: 問題文は社員番号 1 から始まりますが、プログラム内では配列のインデックス(0から開始)に合わせて調整する必要があります(例:上司 \(P_i\) から 1 引く)。

  • 初期値の設定: child_min の初期値は非常に大きな値にするか、最初の部下を処理する際に直接代入するようにします。今回のコードでは -1 をフラグとして使い、最初の部下かどうかを判定しています。

    ソースコード

import sys

def solve():
    # Read all input at once and split by whitespace to handle large input efficiently.
    try:
        input_data = sys.stdin.read().split()
    except EOFError:
        return
        
    if not input_data:
        return
    
    N = int(input_data[0])
    if N < 2:
        # If there's only one employee, they have no subordinates.
        # The unbalance degree of any employee with 1 or fewer subordinates is 0.
        print(0)
        return
        
    # P_2, P_3, ..., P_N are the direct supervisors of employees 2, 3, ..., N.
    # The input provides these N-1 values starting from the second element.
    # P[0] corresponds to P_2, P[1] to P_3, and so on.
    P = list(map(int, input_data[1:N]))
    
    # W_1, W_2, ..., W_N are the workloads assigned to each employee.
    # These N values start from index N in the input_data list.
    W = list(map(int, input_data[N:2*N]))
    
    # S[i] will store the total team workload (subtree sum) of employee i+1.
    # Initially, it is just the workload of the employee themselves.
    S = W
    
    # We need to track the maximum and minimum team workloads among the direct
    # subordinates of each employee to calculate their "unbalance degree".
    child_max = [-1] * N
    child_min = [-1] * N
    child_count = [0] * N
    
    # The problem guarantees that for each employee i (2 <= i <= N), P_i < i.
    # This property allows us to process the tree bottom-up by iterating from
    # N down to 2. When we process employee i, we can be certain that their
    # team workload S[i-1] has already incorporated all workloads from their
    # own subordinates (who must have indices greater than i).
    for i in range(N, 1, -1):
        # p_idx is the 0-indexed ID of the supervisor of employee i.
        # Since P[0] is the supervisor of employee 2, P[i-2] is the supervisor of employee i.
        p_idx = P[i-2] - 1
        
        # s_val is the total team workload of employee i.
        s_val = S[i-1]
        
        # Add employee i's team workload to their supervisor's total team workload.
        S[p_idx] += s_val
        
        # Update the supervisor's records regarding their direct subordinates.
        child_count[p_idx] += 1
        if child_max[p_idx] == -1:
            # This is the first child processed for this supervisor.
            child_max[p_idx] = s_val
            child_min[p_idx] = s_val
        else:
            # Compare this child's team workload with current max and min.
            if s_val > child_max[p_idx]:
                child_max[p_idx] = s_val
            if s_val < child_min[p_idx]:
                child_min[p_idx] = s_val
                
    # The organization's unbalance degree is the maximum unbalance degree among all employees.
    max_unbalance = 0
    for i in range(N):
        # An employee's unbalance degree is only non-zero if they have at least 2 children.
        if child_count[i] >= 2:
            unbalance = child_max[i] - child_min[i]
            if unbalance > max_unbalance:
                max_unbalance = unbalance
                
    # Output the final result.
    sys.stdout.write(str(max_unbalance) + '\n')

if __name__ == '__main__':
    solve()

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

posted:
last update: