公式

C - 水やり / Watering 解説 by admin

GPT 5.2 High

概要

各クエリで「ある頂点 \(X\) の部分木の全頂点に \(D\) を加算」します。木をオイラーツアーで配列に潰し、部分木を区間に変換して区間加算をまとめて処理します。

考察

水やり \((X, D)\) は「\(X\) を根とする部分木の全ノードに一様に \(D\) を足す」操作です。

  • 素朴にやると、毎回 \(X\) の部分木を DFS/BFS で巡って全頂点に \(D\) を加える必要があります。
    すると最悪で 1 回のクエリが \(O(N)\)、\(Q\) 回で \(O(NQ)\) となり、\(N,Q \le 2\times 10^5\) では確実に TLE します。

ここで重要な観察は次の2つです。

  1. 木の部分木は、DFS順(入時刻)で並べると連続区間になる
    つまり、各頂点 \(v\) に「DFSで最初に訪れた時刻 \(tin[v]\)」と「部分木の最後の時刻 \(tout[v]\)」を付けると、
    \(v\) の部分木は配列上の区間 \([tin[v], tout[v]]\) に対応します。

  2. 「区間に一様加算」をたくさん処理するには差分配列が有効
    配列 \(A\) に対し区間 \([l,r]\) に \(D\) を足すなら、差分配列 \(diff\) を使って
    \(diff[l] += D,\ diff[r+1] -= D\) としておき、最後に前から累積和を取れば \(A\) を復元できます。

この2つを組み合わせると、各クエリを \(O(1)\) で記録し、最後に \(O(N)\) で全頂点の答えが出ます。

(簡単な例)
DFS順の配列が [1,2,4,5,3] で、\(X=2\) の部分木が 2,4,5 なら区間は \([1,3]\)。
よって「部分木に加算」は「配列の区間に加算」に置き換えられます。

アルゴリズム

  1. 入力の親配列から子リスト children を作る(根は 1)。
  2. 根 1 から DFS を行い、以下を計算する(オイラーツアーの前順)
    • tin[v]: 訪問順のインデックス
    • tout[v]: 部分木の最後のインデックス
    • euler[t]: 時刻 \(t\) にいる頂点番号(後で配列→頂点へ戻すため)
  3. 差分配列 diff(長さ \(N+2\))を用意し、各クエリ \((x,d)\) について
    • \(l=tin[x],\ r=tout[x]\)
    • diff[l] += d
    • diff[r+1] -= d
  4. t=0..N-1 で累積和 acc を取り、配列上の値を得る:
    • acc += diff[t]
    • その値が euler[t] 番の頂点の最終水分量なので ans[euler[t]] = acc
  5. ans[1..N] を出力する。

計算量

  • 時間計算量: \(O(N + Q)\)(DFSで \(O(N)\)、クエリ処理 \(O(Q)\)、累積和 \(O(N)\))
  • 空間計算量: \(O(N)\)(子リスト、オイラーツアー配列、差分配列など)

実装のポイント

  • 再帰DFSは避ける:\(N=2\times 10^5\) なので Python では再帰が深くなりがちです。コードではスタックを使った反復 DFS で tin/tout を計算しています。

  • tout[u] は「部分木の最後の時刻」なので、DFSで子を全部処理し終わったタイミングで time-1 を入れます(コードの tout[u] = time - 1)。

  • 差分配列で diff[r+1] に触れるため、サイズは N+2 にして境界を安全にしています。

    ソースコード

import sys

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    if not data:
        return
    it = iter(data)
    N = next(it)
    Q = next(it)

    children = [[] for _ in range(N + 1)]
    for v in range(2, N + 1):
        p = next(it)
        children[p].append(v)

    tin = [0] * (N + 1)
    tout = [0] * (N + 1)
    euler = [0] * N

    time = 0
    tin[1] = 0
    euler[0] = 1
    time = 1
    stack = [(1, 0)]
    while stack:
        u, i = stack[-1]
        if i < len(children[u]):
            v = children[u][i]
            stack[-1] = (u, i + 1)
            tin[v] = time
            euler[time] = v
            time += 1
            stack.append((v, 0))
        else:
            tout[u] = time - 1
            stack.pop()

    diff = [0] * (N + 2)
    for _ in range(Q):
        x = next(it)
        d = next(it)
        l = tin[x]
        r = tout[x]
        diff[l] += d
        diff[r + 1] -= d

    ans = [0] * (N + 1)
    acc = 0
    for t in range(N):
        acc += diff[t]
        ans[euler[t]] = acc

    sys.stdout.write(" ".join(map(str, ans[1:])))

if __name__ == "__main__":
    main()

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

投稿日時:
最終更新: