Nim-ACL

Rerooting DP

Rerooting DP is a technique for computing tree DP answers for every possible root.

A normal tree DP fixes one root and computes values from child subtrees to the parent.
However, many problems ask for the answer when each vertex is considered as the root.

Running DFS from every vertex would be O(n^2), but rerooting computes all answers in about O(n) or O(n log n).

Use cases

Typical use cases include:

Idea

Think of DP values as contributions coming from adjacent directions.

The answer for vertex v is obtained by merging contributions from all neighbors and then applying addRoot.

Child-side contributions are computed by an ordinary bottom-up pass.
Parent-side contributions are computed using prefix/suffix products of neighbor contributions, excluding one child at a time.

This module asks you to provide three functions:

Import

import atcoder/extra/graph/rerooting_dp

Types

type RerootingEdge
type RerootingGraph

Constructor

proc initRerootingGraph(n: int): RerootingGraph
proc addEdge(g: var RerootingGraph, u, v: int): int

Run

proc rerooting[T](
  g: RerootingGraph,
  identity: T,
  merge: proc(a, b: T): T,
  addRoot: proc(x: T, v: int): T,
  addEdge: proc(x: T, edgeId: int, fromVertex: int, toVertex: int): T
): seq[T]

proc rerooting[T](
  g: RerootingGraph,
  identity: T,
  merge: proc(a, b: T): T,
  addRoot: proc(x: T, v: int): T
): seq[T]

If addEdge is omitted, values are unchanged across edges.

rerootingDP is also available as an alias.

Example: sum of distances from each vertex

import std/sequtilsimport atcoder/extra/graph/rerooting_dptype DP = tuple[size, dist: int]var g = initRerootingGraph(4)discard g.addEdge(0, 1)discard g.addEdge(1, 2)discard g.addEdge(2, 3)proc merge(a, b: DP): DP =  (size: a.size + b.size, dist: a.dist + b.dist)proc addRoot(x: DP, v: int): DP =  (size: x.size + 1, dist: x.dist)proc addEdge(x: DP, edgeId: int, fromVertex: int, toVertex: int): DP =  (size: x.size, dist: x.dist + x.size)let ans = g.rerooting((size: 0, dist: 0), merge, addRoot, addEdge)doAssert ans.mapIt(it.dist) == @[6, 4, 4, 6]

Here, size is the number of vertices in the component, and dist is the sum of distances from the current root.
When crossing one edge, all distances in that component increase by one, so dist becomes dist + size.

Example: farthest distance from each vertex

var g = initRerootingGraph(6)discard g.addEdge(0, 1)discard g.addEdge(0, 2)discard g.addEdge(1, 3)discard g.addEdge(1, 4)discard g.addEdge(4, 5)proc merge(a, b: int): int = max(a, b)proc addRoot(x: int, v: int): int = xproc addEdge(x: int, edgeId: int, fromVertex: int, toVertex: int): int = x + 1let ans = g.rerootingDP(0, merge, addRoot, addEdge)doAssert ans == @[3, 2, 4, 3, 3, 4]

Nim-friendly usage

Write merge, addRoot, and addEdge as ordinary procs and call with UFCS.

let ans = g.rerooting(identity, merge, addRoot, addEdge)

If edge weights are needed, use edgeId inside addEdge to refer to an external array.

proc addEdge(x: DP, edgeId: int, fromVertex: int, toVertex: int): DP =
  let w = weight[edgeId]
  ...

Complexity

For n vertices:

Notes