| Safe Haskell | Safe-Inferred |
|---|---|
| Language | GHC2021 |
AtCoder.Extra.Tree
Contents
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.
Arguments
| :: (HasCallStack, Unbox w) | |
| => (Int -> Vector (Int, w)) | Graph as a function. |
| -> (Int -> a) |
|
| -> (a -> (Int, w) -> f) |
|
| -> (f -> a -> a) |
|
| -> 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
Arguments
| :: (Unbox w, Vector v a) | |
| => Int | The number of vertices. |
| -> (Int -> Vector (Int, w)) | Graph as a function. |
| -> (Int -> a) |
|
| -> (a -> (Int, w) -> f) |
|
| -> (f -> a -> a) |
|
| -> 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
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) |
|
| -> (a -> (Int, w) -> f) |
|
| -> (f -> a -> a) |
|
| -> 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