公式

C - 水やり / Watering 解説 by admin

Claude 4.6 Opus (Thinking)

概要

根付き木において、指定された頂点の部分木全体に値を加算するクエリを複数回処理し、最終的な各頂点の値を求める問題です。オイラーツアーと差分配列(いもす法)を組み合わせることで効率的に解けます。

考察

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

各クエリで植物 \(X_j\) を根とする部分木のすべての頂点に \(D_j\) を加算する、という操作を愚直に行うと、最悪の場合1回のクエリで \(O(N)\) 個の頂点を更新する必要があります。クエリが \(Q\) 回あるので、全体で \(O(NQ)\) となり、\(N, Q\) がともに最大 \(2 \times 10^5\) の場合は \(4 \times 10^{10}\) 回の操作となり TLE(時間超過) になります。

重要な気づき:部分木は連続区間に対応する

根付き木に対して オイラーツアー(DFS順の訪問時刻) を求めると、任意の頂点 \(v\) の部分木に属する頂点は、オイラーツアー上で 連続した区間 \([\text{in}[v], \text{out}[v]]\) に対応します。

これにより、「部分木全体に値を加算する」操作は「配列の連続区間に値を加算する」操作に変換できます。連続区間への加算は 差分配列(いもす法) を使えば \(O(1)\) で処理できます。

具体例

例えば以下のような木を考えます(\(N = 5\)、根は \(1\)):

    1
   / \
  2   3
 / \
4   5

オイラーツアーで訪問順に番号を振ると:

頂点 in_time out_time
1 0 4
2 1 3
3 4 4
4 2 2
5 3 3

頂点 \(2\) の部分木は \(\{2, 4, 5\}\) で、オイラーツアー上では区間 \([1, 3]\) に対応します。この区間に \(D\) を加算すれば、部分木全体への加算が実現できます。

アルゴリズム

  1. 木の構築: 入力から親子関係を読み取り、隣接リスト(子のリスト)を構築する。

  2. オイラーツアー: 根(頂点 \(1\))からDFSを行い、各頂点の訪問開始時刻 \(\text{in}[v]\) と終了時刻 \(\text{out}[v]\) を記録する。頂点 \(v\) の部分木はオイラーツアー上で区間 \([\text{in}[v], \text{out}[v]]\) に対応する。

  3. 差分配列による区間加算: サイズ \(N\) の差分配列 \(\text{diff}\) を用意する。各クエリ \((X_j, D_j)\) に対して:

    • \(\text{diff}[\text{in}[X_j]]\) に \(D_j\) を加算
    • \(\text{diff}[\text{out}[X_j] + 1]\) から \(D_j\) を減算(範囲内の場合)
  4. 累積和の計算: 差分配列の累積和を取ることで、オイラーツアー上の各位置における最終的な水分量を得る。

  5. 結果の復元: 各頂点 \(v\) の水分量は、累積和配列の \(\text{in}[v]\) 番目の値である。

計算量

  • 時間計算量: \(O(N + Q)\)
    • オイラーツアー: \(O(N)\)
    • 各クエリの処理: \(O(1)\) × \(Q\) 回 = \(O(Q)\)
    • 累積和と結果の復元: \(O(N)\)
  • 空間計算量: \(O(N)\)

実装のポイント

  • 再帰DFSの回避: \(N\) が最大 \(2 \times 10^5\) のため、Pythonの再帰上限に引っかかる可能性がある。スタックを用いた反復DFSでオイラーツアーを実装している。

  • 差分配列の境界処理: \(\text{out}[X_j] + 1\) が配列の範囲外(\(\geq N\))の場合は減算をスキップする必要がある。これは部分木が配列の末尾まで続く場合に該当する。

  • 高速入力: sys.stdin.buffer.read() で一括読み込みすることで、Pythonでも入出力がボトルネックにならないようにしている。

    ソースコード

import sys
from collections import deque

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
    
    parent = [0] * (N + 1)
    children = [[] for _ in range(N + 1)]
    
    for i in range(2, N + 1):
        p = int(input_data[idx]); idx += 1
        parent[i] = p
        children[p].append(i)
    
    # Euler tour to get in/out times
    in_time = [0] * (N + 1)
    out_time = [0] * (N + 1)
    
    # Iterative DFS for Euler tour
    timer = 0
    stack = [(1, False)]
    while stack:
        node, visited = stack.pop()
        if visited:
            out_time[node] = timer - 1
        else:
            in_time[node] = timer
            timer += 1
            stack.append((node, True))
            for c in reversed(children[node]):
                stack.append((c, False))
    
    # Use difference array on Euler tour order
    # For subtree of X: add D to all positions in [in_time[X], out_time[X]]
    diff = [0] * (N + 2)
    
    for _ in range(Q):
        X = int(input_data[idx]); idx += 1
        D = int(input_data[idx]); idx += 1
        l = in_time[X]
        r = out_time[X]
        diff[l] += D
        if r + 1 <= N - 1:
            diff[r + 1] -= D
    
    # Prefix sum
    result = [0] * N
    current = 0
    for i in range(N):
        current += diff[i]
        result[i] = current
    
    # Map back: node with in_time[node] = i gets result[i]
    ans = [0] * (N + 1)
    for node in range(1, N + 1):
        ans[node] = result[in_time[node]]
    
    sys.stdout.write(' '.join(map(str, ans[1:])) + '\n')

main()

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

投稿日時:
最終更新: