Official

E - 会社の組織と給与 / Company Organization and Salaries Editorial by admin

Claude 4.6 Opus (Thinking)

概要

根付き木で表される会社の組織において、社員の給与更新と、ある社員の部下のうちその社員より給与が高い人数を数えるクエリを効率的に処理する問題です。

考察

部分木を区間に変換する

まず「社員 \(v\) のすべての部下」を高速に扱うために、木構造を配列上の連続区間に変換したいと考えます。これはオイラーツアー(Euler Tour)を用いることで実現できます。

DFS の訪問順に社員を並べると、社員 \(v\) の部分木に属する社員は、配列上で連続した区間 \([\mathrm{tin}[v], \mathrm{tout}[v]]\) に対応します。部下(\(v\) 自身を除く)は区間 \([\mathrm{tin}[v]+1, \mathrm{tout}[v]]\) です。

素朴なアプローチの問題点

クエリ2で毎回部下全員を走査すると、最悪 \(O(N)\) かかり、\(Q\) 個のクエリで \(O(NQ)\) となります。\(N, Q \leq 10^5\) のとき最大 \(10^{10}\) 回の操作となり、TLEします。

平方分割で解決

配列上の区間に対する「ある値より大きい要素の個数」を求める問題に帰着できたので、平方分割(Sqrt Decomposition)を適用します。

配列をサイズ \(B \approx \sqrt{N}\) ごとのブロックに分割し、各ブロック内の要素をソート済みリストとしても保持します。こうすることで:

  • 更新:対応するブロックのソート済みリストから旧値を削除し、新値を挿入 → \(O(B)\)
  • 区間クエリ:完全に含まれるブロックは二分探索で \(O(\log B)\)、端のブロックは線形走査で \(O(B)\) → 全体 \(O(\frac{N}{B} \log B + B)\)

\(B = \sqrt{N}\) とすれば、各クエリが \(O(\sqrt{N} \log N)\) 程度で処理できます。

アルゴリズム

  1. オイラーツアーの構築:社員 \(1\) を根として反復的DFSを行い、各社員の訪問開始時刻 \(\mathrm{tin}[v]\) と終了時刻 \(\mathrm{tout}[v]\) を記録する。DFS順に給与を並べた配列 euler_salary を作る。

  2. 平方分割の初期化:euler_salary をブロックサイズ \(B = 250\) 程度で分割し、各ブロック内の値をソートしたリスト blocks[b] を作成する。

  3. クエリ処理:

    • 更新クエリ 1 v x:社員 \(v\) のオイラーツアー上の位置を求め、euler_salary の値を更新。対応するブロックのソート済みリストから旧値を削除し新値を挿入する。
    • 調査クエリ 2 v:部下の区間 \([l, r] = [\mathrm{tin}[v]+1, \mathrm{tout}[v]]\) について、\(A[v]\) より厳密に大きい要素数を数える。完全ブロックには bisect_right で二分探索し、端ブロックは線形走査する。

計算量

  • 時間計算量: \(O(N \log N + Q \cdot (\sqrt{N} \log N + \sqrt{N}))\)(各クエリ \(O(\sqrt{N} \log N)\) 程度。更新は \(O(\sqrt{N})\))
  • 空間計算量: \(O(N)\)

実装のポイント

  • オイラーツアーを反復DFSで実装:\(N \leq 10^5\) では再帰DFSだとPythonのスタック制限に引っかかるため、明示的なスタックを使って反復的に処理する。

  • ソート済みリストの更新:bisect_right(block, old_val - 1) で旧値の位置を見つけて pop し、insort で新値を挿入する。これにより \(O(B)\) で更新できる。

  • 部下が0人の場合:\(l > r\) となるケースは即座に 0 を出力する(葉ノードなど)。

  • ブロックサイズの選択:\(B = 250\) 程度が実測で良好。Pythonでは定数倍が重いため、完全ブロック部分で二分探索を使って高速化することが重要。

    ソースコード

import sys
from bisect import bisect_right, insort

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    Q = int(input_data[idx]); idx += 1
    
    A = [0] * (N + 1)
    for i in range(1, N + 1):
        A[i] = int(input_data[idx]); idx += 1
    
    children = [[] for _ in range(N + 1)]
    for i in range(2, N + 1):
        p = int(input_data[idx]); idx += 1
        children[p].append(i)
    
    # Euler tour to get subtree ranges
    tin = [0] * (N + 1)
    tout = [0] * (N + 1)
    order = []  # euler tour order
    
    # Iterative DFS for Euler tour
    timer = 0
    stack = [(1, False)]
    while stack:
        node, visited = stack.pop()
        if visited:
            tout[node] = timer - 1
        else:
            tin[node] = timer
            timer += 1
            order.append(node)
            stack.append((node, True))
            for c in reversed(children[node]):
                stack.append((c, False))
    
    # For N, Q <= 5*10^4, we can use sqrt decomposition
    # Block size for sqrt decomposition on the Euler tour array
    # Each block maintains a sorted list of salary values
    
    BLOCK = 250  # sqrt(50000) ~ 224, use 250
    num_blocks = (N + BLOCK - 1) // BLOCK
    
    # euler[i] = salary of the node at position i in euler tour
    euler_salary = [0] * N
    euler_node = [0] * N  # which node is at position i
    node_pos = [0] * (N + 1)  # position of node in euler tour
    
    for i in range(N):
        node = order[i]
        euler_salary[i] = A[node]
        euler_node[i] = node
        node_pos[node] = i
    
    # Build sorted blocks
    blocks = []
    for b in range(num_blocks):
        start = b * BLOCK
        end = min(start + BLOCK, N)
        sorted_block = sorted(euler_salary[start:end])
        blocks.append(sorted_block)
    
    def update(node, new_val):
        pos = node_pos[node]
        old_val = euler_salary[pos]
        euler_salary[pos] = new_val
        A[node] = new_val
        b = pos // BLOCK
        # Remove old_val and insert new_val in sorted block
        block = blocks[b]
        # Find and remove old_val
        lo = bisect_right(block, old_val - 1)
        # block[lo] should be old_val
        # Use index to find exact position
        block.pop(lo)
        insort(block, new_val)
    
    def count_greater(l, r, val):
        # Count elements strictly greater than val in euler_salary[l..r]
        if l > r:
            return 0
        result = 0
        lb = l // BLOCK
        rb = r // BLOCK
        
        if lb == rb:
            # Same block, brute force
            for i in range(l, r + 1):
                if euler_salary[i] > val:
                    result += 1
            return result
        
        # Left partial block
        left_end = (lb + 1) * BLOCK - 1
        for i in range(l, left_end + 1):
            if euler_salary[i] > val:
                result += 1
        
        # Full blocks
        for b in range(lb + 1, rb):
            block = blocks[b]
            # Count elements > val
            pos = bisect_right(block, val)
            result += len(block) - pos
        
        # Right partial block
        right_start = rb * BLOCK
        for i in range(right_start, r + 1):
            if euler_salary[i] > val:
                result += 1
        
        return result
    
    out = []
    for _ in range(Q):
        t = int(input_data[idx]); idx += 1
        if t == 1:
            v = int(input_data[idx]); idx += 1
            x = int(input_data[idx]); idx += 1
            update(v, x)
        else:
            v = int(input_data[idx]); idx += 1
            # Subtree of v: positions tin[v] to tout[v]
            # Descendants only: tin[v]+1 to tout[v]
            l = tin[v] + 1
            r = tout[v]
            if l > r:
                out.append('0')
            else:
                val = A[v]
                ans = count_greater(l, r, val)
                out.append(str(ans))
    
    sys.stdout.write('\n'.join(out) + '\n')

main()

この解説は claude4.6opus-thinking によって生成されました。