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つです。
木の部分木は、DFS順(入時刻)で並べると連続区間になる
つまり、各頂点 \(v\) に「DFSで最初に訪れた時刻 \(tin[v]\)」と「部分木の最後の時刻 \(tout[v]\)」を付けると、
\(v\) の部分木は配列上の区間 \([tin[v], tout[v]]\) に対応します。「区間に一様加算」をたくさん処理するには差分配列が有効
配列 \(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]\)。
よって「部分木に加算」は「配列の区間に加算」に置き換えられます。
アルゴリズム
- 入力の親配列から子リスト
childrenを作る(根は 1)。 - 根 1 から DFS を行い、以下を計算する(オイラーツアーの前順)
tin[v]: 訪問順のインデックス
tout[v]: 部分木の最後のインデックス
euler[t]: 時刻 \(t\) にいる頂点番号(後で配列→頂点へ戻すため)
- 差分配列
diff(長さ \(N+2\))を用意し、各クエリ \((x,d)\) について
- \(l=tin[x],\ r=tout[x]\)
diff[l] += d
diff[r+1] -= d
- \(l=tin[x],\ r=tout[x]\)
t=0..N-1で累積和accを取り、配列上の値を得る:
acc += diff[t]
- その値が
euler[t]番の頂点の最終水分量なのでans[euler[t]] = acc
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 によって生成されました。
投稿日時:
最終更新: