全方位木DPで根を付け替える

各頂点を根にして解き直すと O(N²)。根を隣へ動かしたとき、答えがどう変わるかを考えます。

言語：Python / 計算量：O(N)

前提：DFSと木の親子関係

考え方

1. 1つの根で部分木サイズと距離和を求めます。
2. 親vから子uへ根を動かすと、u側の頂点は1近くなります。
3. 残りの頂点は1遠くなるので ans[u]=ans[v]+N-2*size[u]。

具体例

3頂点の鎖0—1—2。0からの距離和3。size[1]=2なので1の答えは3+3-4=2です。

実装

n=3;g=[[1],[0,2],[1]]
parent=[-1]*n;parent[0]=0;order=[0];depth=[0]*n
for v in order:
    for u in g[v]:
        if u==parent[v]:continue
        parent[u]=v;depth[u]=depth[v]+1;order.append(u)
size=[1]*n
for v in order[:0:-1]:size[parent[v]]+=size[v]
ans=[0]*n;ans[0]=sum(depth)
for v in order[1:]:ans[v]=ans[parent[v]]+n-2*size[v]
print(ans)

注意する条件

木であることが前提です。一般グラフでは親以外の訪問済み頂点もあり、この走査は使えません。

確認問題

3頂点の鎖の中央から全頂点への距離和は？

実装課題

講義の木を4頂点の鎖0—1—2—3に変え、各頂点の距離和を空白区切りで出す。入力なし。

出力:
6 4 4 6

n=4;g=[[1],[0,2],[1,3],[2]]
p=[-1]*n;p[0]=0;o=[0];d=[0]*n
for v in o:
    for u in g[v]:
        if u==p[v]:continue
        p[u]=v;d[u]=d[v]+1;o.append(u)
s=[1]*n
for v in o[:0:-1]:s[p[v]]+=s[v]
a=[0]*n;a[0]=sum(d)
for v in o[1:]:a[v]=a[p[v]]+n-2*s[v]
print(*a)

読了の記録・下書き・メモへ