rerooting_dp

ソースコード

from titan_pylib.graph.rerooting_dp import rerooting_dp

view on github

展開済みコード

  1# from titan_pylib.graph.rerooting_dp import rerooting_dp
  2from typing import Callable, TypeVar
  3
  4T = TypeVar("T")
  5E = TypeVar("E")
  6
  7
  8def rerooting_dp(
  9    G: list[list[tuple[int, int]]],
 10    merge: Callable[[T, T], T],
 11    apply_vertex: Callable[[T, int], T],
 12    apply_edge: Callable[[T, E, int, int], T],
 13    e: T,
 14) -> None:
 15    """全方位木DP
 16
 17    Args:
 18        G (list[list[tuple[int, int]]]):
 19        merge (Callable[[T, T], T]):
 20        apply_vertex (Callable[[T, int], T]):
 21        apply_edge (Callable[[T, E, int, int], T]):
 22        e (T):
 23
 24    Returns:
 25        _type_:
 26    """
 27    n = len(G)
 28
 29    dp: list[list[T]] = [[e] * len(g) for g in G]
 30    ans = [e] * len(G)
 31
 32    root = 0
 33    par = [-2] * n
 34    par[root] = -1
 35    toposo = [root]
 36
 37    todo = [root]
 38    while todo:
 39        v = todo.pop()
 40        for x, _ in G[v]:
 41            if par[x] != -2:
 42                continue
 43            par[x] = v
 44            toposo.append(x)
 45            todo.append(x)
 46
 47    arr = [e] * n
 48    for v in toposo[::-1]:
 49        dp_v = dp[v]
 50        acc = e
 51        for i, (x, c) in enumerate(G[v]):
 52            if x == par[v]:
 53                continue
 54            dp_v[i] = apply_edge(arr[x], c, x, v)
 55            acc = merge(acc, dp_v[i])
 56        arr[v] = apply_vertex(acc, v)
 57
 58    dp_par = [e] * n
 59    acc_l = [e] * (n + 1)
 60    acc_r = [e] * (n + 1)
 61    for v in toposo:
 62        dp_v = dp[v]
 63        for i, (x, _) in enumerate(G[v]):
 64            if x == par[v]:
 65                dp_v[i] = dp_par[v]
 66                break
 67        d = len(dp_v)
 68        for i in range(d):
 69            acc_l[i + 1] = merge(acc_l[i], dp_v[i])
 70            acc_r[i + 1] = merge(acc_r[i], dp_v[d - i - 1])
 71        ans[v] = apply_vertex(acc_l[d], v)
 72        for i, (x, c) in enumerate(G[v]):
 73            if x == par[v]:
 74                continue
 75            dp_par[x] = apply_edge(
 76                apply_vertex(
 77                    merge(acc_l[i], acc_r[d - i - 1]),
 78                    v,
 79                ),
 80                c,
 81                v,
 82                x,
 83            )
 84    return ans
 85
 86
 87# apply_vertex(dp_x: T, v: int) -> T:
 88#        v          } return
 89#  --------------   }
 90# |  /   |   \  |   }
 91# | o    o    o | } dp_x (mergeしたもの)
 92# | △   △   △ |
 93#  --------------
 94
 95# apply_edge(dp_x: T, e: E, x: int, v: int) -> T:
 96#   v      } return
 97#   | } e
 98#   x | } dp_x
 99#   △|
100
101# def merge(s: T, t: T) -> T:
102#     """``s`` , ``t`` をマージする"""
103#     ...
104
105# def apply_vertex(dp_x: T, v: int) -> T:
106#     ...
107
108# def apply_edge(dp_x: T, e: E, x: int, v: int) -> T:
109#     ...
110
111# e: T = ...

仕様

rerooting_dp(G: list[list[tuple[int, int]]], merge: Callable[[T, T], T], apply_vertex: Callable[[T, int], T], apply_edge: Callable[[T, E, int, int], T], e: T) → None[source]

全方位木DP

Parameters:

  • G (list[list[tuple[int, int]]])

  • merge (Callable[[T, T], T])

  • apply_vertex (Callable[[T, int], T])

  • apply_edge (Callable[[T, E, int, int], T])

  • e (T)

Return type:

_type_