Safe Haskell	Safe-Inferred
Language	GHC2021

AtCoder.Extra.Tree

Contents

Tree properties
Minimum spanning tree
Tree folding

Description

Generic tree functions.

Since: 1.1.0.0

Synopsis

diameter :: (HasCallStack, Unbox w, Num w, Ord w) => Int -> (Int -> Vector (Int, w)) -> w -> ((Int, Int), w)
diameterPath :: (HasCallStack, Show w, Unbox w, Num w, Ord w) => Int -> (Int -> Vector (Int, w)) -> w -> (Vector Int, w)
mst :: (Num w, Ord w, Unbox w) => Int -> Vector (Int, Int, w) -> (w, Vector Bit, Csr w)
mstBy :: (Num w, Ord w, Unbox w) => (w -> w -> Ordering) -> Int -> Vector (Int, Int, w) -> (w, Vector Bit, Csr w)
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 properties

diameter Source #

Arguments

:: (HasCallStack, Unbox w, Num w, Ord w)
=> Int	The number of vertices.
-> (Int -> Vector (Int, w))	Graph given as a function.
-> w	Distances assigned to unreachable vertices.
-> ((Int, Int), w)	Tuple of (endpoints of the longest path in a tree, distance of it).

\(O(n + m)\) Returns the endpoints of the diameter of a tree and their distance: \(((u, v), w)\).

Example

Expand

>>> import AtCoder.Extra.Graph qualified as Gr
>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Vector.Unboxed qualified as VU
>>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 10), (1, 3, 10)]
>>> let gr = Gr.build 4 $ Gr.swapDupe es
>>> Tree.diameter 4 (Gr.adjW gr) (-1)
((2,3),20)

Since: 1.2.4.0

diameterPath Source #

Arguments

:: (HasCallStack, Show w, Unbox w, Num w, Ord w)
=> Int	The number of vertices.
-> (Int -> Vector (Int, w))	Graph given as a function.
-> w	Distances assigned to unreachable vertices.
-> (Vector Int, w)	Tuple of (the longest path, distance of it).

\(O(n + m)\) Returns the path longest path in a tree and the distance of it.

Example

Expand

>>> import AtCoder.Extra.Graph qualified as Gr
>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Vector.Unboxed qualified as VU
>>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 10), (1, 3, 10)]
>>> let gr = Gr.build 4 $ Gr.swapDupe es
>>> Tree.diameterPath 4 (Gr.adjW gr) (-1)
([2,1,3],20)

Since: 1.2.4.0

Minimum spanning tree

mst :: (Num w, Ord w, Unbox w) => Int -> Vector (Int, Int, w) -> (w, Vector Bit, Csr w) Source #

\(O(m \log m)\) Kruscal's algorithm. Returns edge indices for building a minimum spanning tree.

NOTE: The edges should not be duplicated: only one of \((u, v, w)\) or \((v, u w)\) is required for each edge.

Example

Expand

Create a minimum spanning tree:

>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Vector.Unboxed qualified as VU
>>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 10), (0, 2, 2)]
>>> let (!wSum, !edgeUse, !gr) = Tree.mst 3 es
>>> wSum
3

>>> edgeUse
[1,0,1]

>>> Gr.adj gr 0
[1,2]

Since: 1.2.4.0

mstBy :: (Num w, Ord w, Unbox w) => (w -> w -> Ordering) -> Int -> Vector (Int, Int, w) -> (w, Vector Bit, Csr w) Source #

\(O(m \log m)\) Kruscal's algorithm. Returns edge indices for building a minimum/maximum spanning tree.

NOTE: The edges should not be duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is required for each edge.

Example

Expand

Create a maximum spanning tree:

>>> import AtCoder.Extra.Tree qualified as Tree
>>> import Data.Ord (Down (..))
>>> import Data.Vector.Unboxed qualified as VU
>>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 10), (0, 2, 2)]
>>> let (!wSum, !edgeUse, !gr) = Tree.mstBy (comparing Down) 3 es
>>> wSum
12

>>> edgeUse
[0,1,1]

>>> Gr.adj gr 0
[2]

Since: 1.2.4.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 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.

Constraints

The action monoid \(f\) must be commutative.

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