AtCoder.Extra.Tree

Generic tree functions.

Since: 1.1.0.0

Tree folding

These function are built around the three type parameters: \(w\), \(f\) and \(a\).

• \(w\): Edge weight.
• \(f\): Monoid action to a vertex value. These actions are created from vertex value \(a\) and edge information (Int, w).
• \(a\): Monoid values stored at vertices.

fold

Arguments

:: (HasCallStack, Unbox w)

=> (Int -> Vector (Int, w))

Graph as a function.

-> (Int -> a)

valAt: Assignment of initial vertex values.

-> (a -> (Int, w) -> f)

toF: Converts a vertex value into an action onto a neighbor vertex.

-> (f -> a -> a)

act: Performs an action onto a vertex value.

-> Int

Root vertex.

-> a

Tree folding result from the root vertex.

\(O(n)\) Folds a tree from a root vertex, also known as tree DP.

Example

>>> import AtCoder.Extra.Graph qualified as Gr
>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Semigroup (Sum (..))
>>> import Data.Vector.Unboxed qualified as VU
>>> let gr = Gr.build @(Sum Int) 5 . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)]
>>> type W = Sum Int -- edge weight
>>> type F = Sum Int -- action type
>>> type X = Sum Int -- vertex value
>>> :{
 let res = Tree.fold (gr `Gr.adjW`) valAt toF act 2
       where
         valAt :: Int -> X
         valAt = const $ mempty @(Sum Int)
         toF :: X -> (Int, W) -> F
         toF x (!_i, !dx) = x + dx
         act :: F -> X -> X
         act dx x = dx + x
  in getSum res
:}
4

Since: 1.1.0.0

scan

Arguments

:: (Unbox w, Vector v a)

=> Int

The number of vertices.

-> (Int -> Vector (Int, w))

Graph as a function.

-> (Int -> a)

valAt: Assignment of initial vertex values.

-> (a -> (Int, w) -> f)

toF: Converts a vertex value into an action onto a neighbor vertex.

-> (f -> a -> a)

act: Performs an action onto a vertex value.

-> Int

Root vertex.

-> v a

Tree scanning result from a root vertex.

\(O(n)\) Folds a tree from a root vertex, also known as tree DP. The calculation process on every vertex is recoreded and returned as a vector.

Example

>>> import AtCoder.Extra.Graph qualified as Gr
>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Semigroup (Sum (..))
>>> import Data.Vector.Unboxed qualified as VU
>>> let n = 5
>>> let gr = Gr.build @(Sum Int) n . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)]
>>> type W = Sum Int -- edge weight
>>> type F = Sum Int -- action type
>>> type X = Sum Int -- vertex value
>>> :{
 let res = Tree.scan n (gr `Gr.adjW`) valAt toF act 2
       where
         valAt :: Int -> X
         valAt = const $ mempty @(Sum Int)
         toF :: X -> (Int, W) -> F
         toF x (!_i, !dx) = x + dx
         act :: F -> X -> X
         act dx x = dx + x
  in VU.map getSum res
:}
[0,1,4,1,0]

Since: 1.1.0.0

foldReroot

Arguments

:: forall w f a. (HasCallStack, Unbox w, Unbox a, Unbox f, Monoid f)

=> Int

The number of vertices.

-> (Int -> Vector (Int, w))

Graph as a function.

-> (Int -> a)

valAt:Assignment of initial vertex values.

-> (a -> (Int, w) -> f)

toF: Converts a vertex value into an action onto a neighbor vertex.

-> (f -> a -> a)

act: Performs an action onto a vertex value.

-> Vector a

Tree folding result from every vertex as a root.

\(O(n)\) Folds a tree from every vertex, using the rerooting technique.

Example

>>> import AtCoder.Extra.Graph qualified as Gr
>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Semigroup (Sum (..))
>>> import Data.Vector.Unboxed qualified as VU
>>> let n = 5
>>> let gr = Gr.build @(Sum Int) n . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)]
>>> type W = Sum Int -- edge weight
>>> type F = Sum Int -- action type
>>> type X = Sum Int -- vertex value
>>> :{
 let res = Tree.foldReroot n (gr `Gr.adjW`) valAt toF act
       where
         valAt :: Int -> X
         valAt = const $ mempty @(Sum Int)
         toF :: X -> (Int, W) -> F
         toF x (!_i, !dx) = x + dx
         act :: F -> X -> X
         act dx x = dx + x
  in VU.map getSum res
:}
[4,4,4,4,4]

Since: 1.1.0.0