CodeForces Round #627 Div3 F MaximumWhiteSubtree

内容

CFのRound #627 Div3のF問題、「MaximumWhiteSubtree」をpython(pypy)で解いた。

問題

入力

• 木が与えられる。
• 各ノードは、白か黒に塗られている。

出力

• 各頂点 v について、以下の問題を解く。
  • vを含む　部分木 Sを、[Sに含まれる白の頂点の数 - 黒の頂点の数]が最大になるように求める。

codeforces.com

考えたこと

• 頂点vについての値は、vを木の根として、木全体をDFSとかで辿ってDPすれば、値は求まる。

  • しかし、それを全頂点に対して行うのは間に合わない。

• 全方位木DP で効率的に計算する。

　全方位木DP

• 全ての頂点に対して、その頂点を根とした時の何らかの計算をするときに使える。
  • 今回は、全ての頂点に対して、その頂点を根とした時の、[白頂点の数- 黒頂点の数]の最大値を求めたい。
• 全方位木DPのめっちゃわかりやすい記事 ei1333.hateblo.jp

実装

import sys
input = sys.stdin.readline

N = int(input())
# 各頂点の色
Color = list(map(int, input().split()))
# 各頂点の隣接する頂点
Adj = [[] for _ in range(N)]

for _ in range(N-1):
    a, b = map(lambda x: int(x)-1, input().split())
    Adj[a].append(b)
    Adj[b].append(a)

# 各頂点の親を記録(root=0とした時の親)
Parent = [0]*N
Parent[0] = -1
# 各頂点の根からの距離(root=0とした時)
depth = [0]*N

# rootからBFSする(root=0の時の、各頂点の親を記録していく)
Q = [0]
visited = {0}
d = 1

# BFS
while Q:
    P = []  # next iteration
    for i in Q:
        for j in Adj[i]:
            if j in visited:
                continue
            depth[j] = d
            Parent[j] = i
            visited.add(j)
            P.append(j)

    Q = P
    d += 1

# ここから、root=0の時のscore (white-black)を求めていく。

# depth_kは、[rootからの距離,頂点のid]をもつ。
depth_k = [[depth[i], i] for i in range(N)]

# depthでsortする。
depth_k.sort()

#　答えを入れる(Ans[i]は頂点iをrootとした時のscore(white-black)
Ans = [0]*N

#　root=0の時の、葉から根に向かってscoreを求めていく
for i in reversed(range(1, N)):
    k = depth_k[i][1]
    if Color[k] == 0:  # 頂点kが黒なら、-1する
        Ans[k] -= 1
    else:
        Ans[k] += 1

    # 親の値は、子の値が正ならその子を使って、負なら使わない。
    Ans[Parent[k]] += max(0, Ans[k])

# rootの文を計算する。
if Color[0] == 0:
    Ans[0] -= 1
else:
    Ans[0] += 1

# ここまでで、root=0とした時のscore(white-black)が求められている。

#　ここから、他の頂点を根にした時のscoreを求めていく
for i in range(1, N):
    k = depth_k[i][1]
    Ans[k] = max(Ans[k], min(Ans[k], 0)+Ans[Parent[k]])

print(*Ans)

思ったこと

• 難しい。次出たときも時間内に解けるとは思えないが、使えると強いから覚えておきたい。
• pythonだと通らないし、pypyでもそんなに余裕がないので、c++に乗り換えるべきかなーと最近思い始めた。