Safe Haskell	Safe-Inferred
Language	GHC2021

AtCoder.Extra.Tree

Contents

Tree folding

Description

Generic tree functions.

Since: 1.1.0.0

Synopsis

fold :: (HasCallStack, Unbox w) => (Int -> Vector (Int, w)) -> (Int -> a) -> (a -> (Int, w) -> f) -> (f -> a -> a) -> Int -> a
scan :: (Unbox w, Vector v a) => Int -> (Int -> Vector (Int, w)) -> (Int -> a) -> (a -> (Int, w) -> f) -> (f -> a -> a) -> Int -> v a
foldReroot :: forall w f a. (HasCallStack, Unbox w, Unbox a, Unbox f, Monoid f) => Int -> (Int -> Vector (Int, w)) -> (Int -> a) -> (a -> (Int, w) -> f) -> (f -> a -> a) -> Vector a

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 Source #

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

Expand

>>> 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 Source #

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

Expand

>>> 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 Source #

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

Expand

>>> 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