ABC187 E Through Path

内容

Atcoder 187 E問題(500) Through Pathをpython で　解説AC

問題

atcoder.jp

ポイント

• 木の上で、imos法を使う。
• 適当に頂点を置くと、queryは、次のどちらかに言い換えができる。
  • ある頂点 k の子孫に xを足す
  • ある頂点 kの子孫以外にxを足す
    • rootの子孫にxを足して、kの子孫からxを引く。 *query type2は、a,bをswapすることで、type1と同じになる。

実装

N = int(input())
edges = [list(map(lambda x:int(x)-1, input().split())) for _ in range(N-1)]

Q = int(input())
queries = [list(map(int, input().split())) for _ in range(Q)]

# 隣接リスト
adjes = [[] for _ in range(N)]
for a, b in edges:
    adjes[a].append(b)
    adjes[b].append(a)

# node 0　を　rootとする。
ROOT = 0

# root_ptr[i]: node i の　親の node
root_ptr = [-1]*N
root_ptr[ROOT] = ROOT


# dfs
Q = [ROOT]
while Q:
    cur = Q.pop()
    for nex in adjes[cur]:
        if root_ptr[nex] != -1:
            continue
        root_ptr[nex] = cur
        Q.append(nex)


# queryを捌く
ans = [0]*N
for t, e, x in queries:
    e -= 1
    a, b = edges[e]
    # t=2 and swap a,b <=> t=1
    if t == 2:
        a, b = b, a

    if root_ptr[a] != b:  # a is parent -> bの子孫以外にxを足す
        ans[ROOT] += x  # 木全体
        ans[b] -= x  # bの子孫
    else:  # b is parent -> aの子孫にxを足す
        ans[a] += x  # aの子孫

# imos
Q = [ROOT]
visited = [False]*N
visited[ROOT] = True

# imos on tree
while Q:
    cur = Q.pop()
    for nex in adjes[cur]:
        if visited[nex]:
            continue
        ans[nex] += ans[cur]
        visited[nex] = True
        Q.append(nex)

print("\n".join(map(str, ans)))

終わりに

• おわり
• 公式解説がわかりやすくて良い