全方位木DP

根を 1 つ決めて下から木 DP をしたあと、上から「親の側の部分木の値」を配ることで、すべての頂点を根としたときの値を求めます。再帰は使いません。問題ごとに次の 3 つを決めて渡します。

• merge(a, b): 子から来た値をまとめる演算（結合的かつ可換）
• e: merge の単位元
• add_root(x, v, p): 子の値をまとめた $x$ に頂点 $v$ を加えて、隣の頂点 $p$ へ渡す値を返す。 $p = -1$ のときは、 $v$ を根としたときの答えを返す

関数は次の 1 つです。

• rerooting(g, merge, e, add_root): 木の隣接リスト g[v] = [to, ...] から、各頂点を根としたときの答えのリストを返す。 $O(N)$（merge と add_root が $O(1)$ のとき）

値には整数やタプルのような、書き換えないものを使います。辺の重みを使うときは、add_root の中で辺 $(v, p)$ の重みを足します（重みは (v, p) をキーとする辞書などで引きます）。

例（各頂点から他のすべての頂点への距離の和）:

def merge(a, b):
    return (a[0] + b[0], a[1] + b[1])


def add_root(x, v, p):
    # x = (子の側の頂点数, v からその頂点たちへの距離の和)
    if p == -1:
        return x
    return (x[0] + 1, x[1] + x[0] + 1)


res = rerooting(g, merge, (0, 0), add_root)
# res[v][1] が v から他のすべての頂点への距離の和

コード

def rerooting(g, merge, e, add_root):
    # g[v] = [to, ...]（木）
    # merge(a, b): 子から来た値をまとめる演算（結合的かつ可換）、e: その単位元
    # add_root(x, v, p): 子の値をまとめた x に頂点 v を加えて p へ渡す値（p == -1 なら v の答え）
    n = len(g)
    par = [-1] * n
    order = [0]
    for v in order:
        for to in g[v]:
            if to != par[v]:
                par[to] = v
                order.append(to)
    down = [e] * n
    acc = [e] * n
    for v in reversed(order[1:]):
        down[v] = add_root(acc[v], v, par[v])
        acc[par[v]] = merge(acc[par[v]], down[v])
    up = [e] * n
    res = [e] * n
    for v in order:
        ch = [to for to in g[v] if to != par[v]]
        vals = [down[c] for c in ch]
        if v != 0:
            vals.append(up[v])
        k = len(vals)
        suf = [e] * (k + 1)
        for i in range(k - 1, -1, -1):
            suf[i] = merge(vals[i], suf[i + 1])
        res[v] = add_root(suf[0], v, -1)
        pre = e
        for i, c in enumerate(ch):
            up[c] = add_root(merge(pre, suf[i + 1]), v, c)
            pre = merge(pre, vals[i])
    return res