{-# LANGUAGE RecordWildCards #-}

-- | Heavy-light decomposition is a method for partitioning a tree into segments with consecutive
-- indices. It processes various path queries in \(O(\log n)\) time. For segment tree integration
-- and monoid products, refer to the @TreeMonoid@ module.
--
-- ==== __Overview of the internals__
-- The following is for understanding the internals, not for using the API. Skip to the examples if
-- you want.
--
-- ===== Original tree
--
-- Consider a tree with arbitrary vertex order:
--
-- @
--  0--8--7--3--1--2--12--13--15--14     XX: original Vertex
--     |        |                        --: edge
-- 10--5        11--8--6                 |: edge
--     |
--     4
-- @
--
-- ===== `indexHld`: Vertex -> VertexHld
--
-- The tree vertices are reindexed with `indexHld`, where each segment is assigned consecutive
-- vertex indices:
--
-- @
--  0==1==2==3==4==5==6==7==8==9     XX: VertexHld
--     |        |                    ==: edges on the same semgent
-- 14==13       10==11==12           |: edge between different segments
--     |
--     15
-- @
--
-- Note that vertices on higher (closer to the root) segments are assigned smaller indices. This is
-- very internally very important when calculating `lca`.
--
-- ===== `headHld`: Vertex -> Vertex
--
-- `headHld` points the "head" vertex of each segment. It can be used for finding LCA of two
-- vertices. To find the LCA, move up to the head, go up to the parental segment's vertex and
-- repeat until the two vertices are on the same segment.
--
-- @
--  0==0==0==0==0==0==0==0==0==0     XX: original Vertex
--     |     |
--  5==5     11==11==11
--     |
--     4
-- @
--
-- `headHld` also works for identifying segments. When two vertices are on the same segment, they
-- have the same head.
--
-- ===== `parentHld`: Vertex -> Vertex
--
-- `parentHld` points the parental segment's vertex from a head:
--
-- @
-- (-1)==0==8==7==3==1==2==12==13==15     XX: original Vertex
--       |        |
--    5==8        1==11=8
--       |
--       5
-- @
--
-- ==== __Example__
-- Create an `Hld` for a tree:
--
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0--1--2--3
-- >>> --    +
-- >>> --    +--4--5
-- >>> let n = 6
-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)]
-- >>> let hld = Hld.new tree
--
-- `Hld` can process various queries in \(O(\log n)\) time:
--
-- >>> Hld.ancestor hld 5 3 -- go up three parents from `5`
-- 0
--
-- >>> Hld.jump hld 5 2 3   -- go to the third vertex from `5` to `2`:
-- Just 2
--
-- >>> Hld.lengthBetween hld 5 3 -- get the length (the number of edges) between `5` and `3`:
-- 4
--
-- >>> Hld.path hld 5 3     -- get the path between `5` and `3`:
-- [5,4,1,2,3]
--
-- Our `Hld` is rooted at @0@ vertex and subtree queries are available:
--
-- >>> Hld.isInSubtree hld 2 3 -- `3` is in the subtree of `2`
-- True
--
-- >>> Hld.isInSubtree hld 2 4 -- `4` is not in the subtree of `2`
-- False
--
-- ===== Segment queries
-- Products and segment queries are primarily used by the @TreeMonoid@ module and is not intended
-- for diretct use, but here's some examples. This time the reindex by the HLD is identity:
--
-- >>> Hld.indexHld hld
-- [0,1,2,3,4,5]
--
-- So we can easily understand the outputs:
--
-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnVertices hld 5 3
-- [(5,4),(1,3)]
--
-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnEdges hld 5 3 -- LCA (1) is removed
-- [(5,4),(2,3)]
--
-- >>> Hld.subtreeSegmentInclusive hld 1
-- (1,5)
--
-- @since 1.1.0.0
module AtCoder.Extra.Tree.Hld
  ( -- * Hld
    Hld (..),
    Vertex,
    VertexHld,

    -- * Constructors
    new,
    newAt,

    -- * LCA
    lca,

    -- * Jump
    ancestor,
    jump,

    -- * Path
    lengthBetween,
    path,
    pathSegmentsInclusive,

    -- * Subtree
    subtreeSegmentInclusive,
    isInSubtree,

    -- * Products
    WeightPolicy (..),
    prod,
  )
where

import AtCoder.Extra.Graph qualified as Gr
import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad
import Control.Monad.Fix
import Control.Monad.ST
import Data.Maybe
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)

-- | Original graph vertex.
--
-- @since 1.1.0.0
type Vertex = Int

-- | Vertex reindexed by `indexHld`.
--
-- @since 1.1.0.0
type VertexHld = Vertex

-- | `Hld` partitions a tree into segments and assignes contiguous `VertexHld` for each segment.
--
-- @since 1.1.0.0
data Hld = Hld
  { -- | The root vertex.
    --
    -- @since 1.1.0.0
    Hld -> Vertex
rootHld :: {-# UNPACK #-} !Vertex,
    -- | Maps `Vertex` to the parent `Vertex`. Returns @-1@ for the root node.
    --
    -- @since 1.1.0.0
    Hld -> Vector Vertex
parentHld :: !(VU.Vector Vertex),
    -- | Maps `Vertex` to `VertexHld`, re-indexed vertices contiguous in each segment.
    --
    -- @since 1.1.0.0
    Hld -> Vector Vertex
indexHld :: !(VU.Vector VertexHld),
    -- | Maps `Vertex` to the head `Vertex` of the segment.
    --
    -- @since 1.1.0.0
    Hld -> Vector Vertex
headHld :: !(VU.Vector Vertex),
    -- | Maps `VertexHld` back to `Vertex`. Used for `ancestor` etc.
    --
    -- @since 1.1.0.0
    Hld -> Vector Vertex
revIndexHld :: !(VU.Vector Vertex),
    -- | Maps `Vertex` to their depth from the root. Used for `jump` etc.
    --
    -- @since 1.1.0.0
    Hld -> Vector Vertex
depthHld :: !(VU.Vector Int),
    -- | Maps `Vertex` to the subtree size. This is for subtree products.
    --
    -- @since 1.1.0.0
    Hld -> Vector Vertex
subtreeSizeHld :: !(VU.Vector Int)
  }
  deriving
    ( -- | @since 1.1.0.0
      Vertex -> Hld -> ShowS
[Hld] -> ShowS
Hld -> String
(Vertex -> Hld -> ShowS)
-> (Hld -> String) -> ([Hld] -> ShowS) -> Show Hld
forall a.
(Vertex -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Vertex -> Hld -> ShowS
showsPrec :: Vertex -> Hld -> ShowS
$cshow :: Hld -> String
show :: Hld -> String
$cshowList :: [Hld] -> ShowS
showList :: [Hld] -> ShowS
Show,
      -- | @since 1.1.0.0
      Hld -> Hld -> Bool
(Hld -> Hld -> Bool) -> (Hld -> Hld -> Bool) -> Eq Hld
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Hld -> Hld -> Bool
== :: Hld -> Hld -> Bool
$c/= :: Hld -> Hld -> Bool
/= :: Hld -> Hld -> Bool
Eq
    )

-- | \(O(n)\) Creates an `Hld` with \(0\) as the root vertex.
--
-- @since 1.1.0.0
{-# INLINE new #-}
new :: forall w. (HasCallStack) => Gr.Csr w -> Hld
new :: forall w. HasCallStack => Csr w -> Hld
new Csr w
tree = Csr w -> Vertex -> Hld
forall w. HasCallStack => Csr w -> Vertex -> Hld
newAt Csr w
tree Vertex
0

-- | \(O(n)\) Creates an `Hld` with a root vertex specified.
--
-- @since 1.1.0.0
{-# INLINE newAt #-}
newAt :: forall w. (HasCallStack) => Gr.Csr w -> Vertex -> Hld
newAt :: forall w. HasCallStack => Csr w -> Vertex -> Hld
newAt Csr w
tree Vertex
root = (forall s. ST s Hld) -> Hld
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Hld) -> Hld) -> (forall s. ST s Hld) -> Hld
forall a b. (a -> b) -> a -> b
$ do
  -- Re-create adjacent vertices so that the biggest subtree's head vertex comes first.
  --
  -- We /could/ instead record the biggest adjacent subtree vertex for each vertex, but the other
  -- DFS would be harder.
  let (!Csr w
tree', !Vector Vertex
parent, !Vector Vertex
depths, !Vector Vertex
subtreeSize) = (forall s.
 ST s (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
-> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex)
forall a. (forall s. ST s a) -> a
runST ((forall s.
  ST s (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
 -> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
-> (forall s.
    ST s (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
-> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex)
forall a b. (a -> b) -> a -> b
$ do
        MVector s Vertex
adjVec <- Vector Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw (Csr w -> Vector Vertex
forall w. Csr w -> Vector Vertex
Gr.adjCsr Csr w
tree)
        MVector s Vertex
parent_ <- Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> m (MVector (PrimState m) a)
VUM.unsafeNew Vertex
n
        MVector s Vertex
depths_ <- Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> m (MVector (PrimState m) a)
VUM.unsafeNew Vertex
n
        MVector s Vertex
subtreeSize_ <- Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> m (MVector (PrimState m) a)
VUM.unsafeNew Vertex
n

        Vertex
_ <- (\(Vertex -> Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vertex -> Vertex -> ST s Vertex
f -> ((Vertex -> Vertex -> Vertex -> ST s Vertex)
 -> Vertex -> Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vertex -> Vertex -> ST s Vertex
forall a. (a -> a) -> a
fix (Vertex -> Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vertex -> Vertex -> ST s Vertex
f Vertex
0 (-Vertex
1) Vertex
root) (((Vertex -> Vertex -> Vertex -> ST s Vertex)
  -> Vertex -> Vertex -> Vertex -> ST s Vertex)
 -> ST s Vertex)
-> ((Vertex -> Vertex -> Vertex -> ST s Vertex)
    -> Vertex -> Vertex -> Vertex -> ST s Vertex)
-> ST s Vertex
forall a b. (a -> b) -> a -> b
$ \Vertex -> Vertex -> Vertex -> ST s Vertex
loop Vertex
depth Vertex
p Vertex
v1 -> do
          MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.write MVector s Vertex
MVector (PrimState (ST s)) Vertex
parent_ Vertex
v1 Vertex
p
          MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.write MVector s Vertex
MVector (PrimState (ST s)) Vertex
depths_ Vertex
v1 Vertex
depth

          (!Vertex
size1, !Vertex
eBig) <-
            ((Vertex, Vertex) -> (Vertex, Vertex) -> ST s (Vertex, Vertex))
-> (Vertex, Vertex)
-> Vector (Vertex, Vertex)
-> ST s (Vertex, Vertex)
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
              ( \(!Vertex
size1, !Vertex
eBig) (!Vertex
e2, !Vertex
v2) -> do
                  if Vertex
v2 Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
p
                    then (Vertex, Vertex) -> ST s (Vertex, Vertex)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
size1, Vertex
eBig)
                    else do
                      Vertex
size2 <- Vertex -> Vertex -> Vertex -> ST s Vertex
loop (Vertex
depth Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) Vertex
v1 Vertex
v2
                      -- NOTE: It's `>` because we should swap at least once if there's some vertex other
                      -- that the parent_.
                      (Vertex, Vertex) -> ST s (Vertex, Vertex)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
size1 Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
size2, if Vertex
size1 Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
> Vertex
size2 then Vertex
eBig else Vertex
e2)
              )
              (Vertex
1 :: Int, -Vertex
1)
              (Csr w
tree Csr w -> Vertex -> Vector (Vertex, Vertex)
forall w.
HasCallStack =>
Csr w -> Vertex -> Vector (Vertex, Vertex)
`Gr.eAdj` Vertex
v1)

          -- move the biggest subtree's head to the first adjacent vertex.
          -- it means the "heavy edge" or the longest segment.
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Vertex
eBig Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= -Vertex
1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> Vertex -> m ()
VGM.swap MVector s Vertex
MVector (PrimState (ST s)) Vertex
adjVec Vertex
eBig (Vertex -> ST s ()) -> Vertex -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Vertex, Vertex) -> Vertex
forall a b. (a, b) -> a
fst (Vector (Vertex, Vertex) -> (Vertex, Vertex)
forall (v :: * -> *) a. Vector v a => v a -> a
VG.head (Csr w
tree Csr w -> Vertex -> Vector (Vertex, Vertex)
forall w.
HasCallStack =>
Csr w -> Vertex -> Vector (Vertex, Vertex)
`Gr.eAdj` Vertex
v1))

          -- record subtree size
          MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.write MVector s Vertex
MVector (PrimState (ST s)) Vertex
subtreeSize_ Vertex
v1 Vertex
size1

          Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
size1

        !Vector Vertex
vec <- MVector (PrimState (ST s)) Vertex -> ST s (Vector Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Vertex
MVector (PrimState (ST s)) Vertex
adjVec
        (Csr w
tree {Gr.adjCsr = vec},,,)
          (Vector Vertex
 -> Vector Vertex
 -> Vector Vertex
 -> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
-> ST s (Vector Vertex)
-> ST
     s
     (Vector Vertex
      -> Vector Vertex
      -> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Vertex -> ST s (Vector Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Vertex
MVector (PrimState (ST s)) Vertex
parent_
          ST
  s
  (Vector Vertex
   -> Vector Vertex
   -> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
-> ST s (Vector Vertex)
-> ST
     s
     (Vector Vertex
      -> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Vertex -> ST s (Vector Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Vertex
MVector (PrimState (ST s)) Vertex
depths_
          ST
  s
  (Vector Vertex
   -> (Csr w, Vector Vertex, Vector Vertex, Vector Vertex))
-> ST s (Vector Vertex)
-> ST s (Csr w, Vector Vertex, Vector Vertex, Vector Vertex)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Vertex -> ST s (Vector Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Vertex
MVector (PrimState (ST s)) Vertex
subtreeSize_

  -- vertex -> reindexed vertex index
  MVector s Vertex
indices <- Vertex -> Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
n (-Vertex
1 :: Int)

  -- vertex -> head vertex of the segment
  MVector s Vertex
heads <- Vertex -> Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
n (-Vertex
1 :: Int)

  Vertex
_ <- (\(Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex
f -> ((Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
 -> Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex
forall a. (a -> a) -> a
fix (Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex
f (Vertex
0 :: Int) Vertex
root (-Vertex
1) Vertex
root) (((Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
  -> Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
 -> ST s Vertex)
-> ((Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
    -> Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex)
-> ST s Vertex
forall a b. (a -> b) -> a -> b
$ \Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex
loop Vertex
acc Vertex
h Vertex
p Vertex
v1 -> do
    -- reindex:
    MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.write MVector s Vertex
MVector (PrimState (ST s)) Vertex
indices Vertex
v1 Vertex
acc
    let !acc' :: Vertex
acc' = Vertex
acc Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1

    MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.write MVector s Vertex
MVector (PrimState (ST s)) Vertex
heads Vertex
v1 Vertex
h

    -- when the first vertex is within the same segment:
    let (!Vertex
adj1, !Vector Vertex
rest) = Maybe (Vertex, Vector Vertex) -> (Vertex, Vector Vertex)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Vertex, Vector Vertex) -> (Vertex, Vector Vertex))
-> Maybe (Vertex, Vector Vertex) -> (Vertex, Vector Vertex)
forall a b. (a -> b) -> a -> b
$ Vector Vertex -> Maybe (Vertex, Vector Vertex)
forall a. Unbox a => Vector a -> Maybe (a, Vector a)
VU.uncons (Csr w
tree' Csr w -> Vertex -> Vector Vertex
forall w. HasCallStack => Csr w -> Vertex -> Vector Vertex
`Gr.adj` Vertex
v1)
    Vertex
acc'' <-
      if Vertex
adj1 Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
p
        then Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
acc'
        else Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex
loop Vertex
acc' Vertex
h Vertex
v1 Vertex
adj1

    -- the others are in other segments:
    (Vertex -> Vertex -> ST s Vertex)
-> Vertex -> Vector Vertex -> ST s Vertex
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
      ( \Vertex
a Vertex
v2 -> do
          if Vertex
v2 Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
p
            then Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
a
            else Vertex -> Vertex -> Vertex -> Vertex -> ST s Vertex
loop Vertex
a Vertex
v2 Vertex
v1 Vertex
v2
      )
      Vertex
acc''
      Vector Vertex
rest

  !Vector Vertex
indices' <- MVector (PrimState (ST s)) Vertex -> ST s (Vector Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Vertex
MVector (PrimState (ST s)) Vertex
indices
  let !revIndex :: Vector Vertex
revIndex = Vector Vertex -> Vector (Vertex, Vertex) -> Vector Vertex
forall a. Unbox a => Vector a -> Vector (Vertex, a) -> Vector a
VU.update (Vertex -> Vertex -> Vector Vertex
forall a. Unbox a => Vertex -> a -> Vector a
VU.replicate Vertex
n (-Vertex
1)) (Vector (Vertex, Vertex) -> Vector Vertex)
-> Vector (Vertex, Vertex) -> Vector Vertex
forall a b. (a -> b) -> a -> b
$ (Vertex -> Vertex -> (Vertex, Vertex))
-> Vector Vertex -> Vector (Vertex, Vertex)
forall a b.
(Unbox a, Unbox b) =>
(Vertex -> a -> b) -> Vector a -> Vector b
VU.imap ((Vertex -> Vertex -> (Vertex, Vertex))
-> Vertex -> Vertex -> (Vertex, Vertex)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,)) Vector Vertex
indices'

  Vertex
-> Vector Vertex
-> Vector Vertex
-> Vector Vertex
-> Vector Vertex
-> Vector Vertex
-> Vector Vertex
-> Hld
Hld Vertex
root Vector Vertex
parent Vector Vertex
indices'
    (Vector Vertex
 -> Vector Vertex -> Vector Vertex -> Vector Vertex -> Hld)
-> ST s (Vector Vertex)
-> ST s (Vector Vertex -> Vector Vertex -> Vector Vertex -> Hld)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Vertex -> ST s (Vector Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Vertex
MVector (PrimState (ST s)) Vertex
heads
    ST s (Vector Vertex -> Vector Vertex -> Vector Vertex -> Hld)
-> ST s (Vector Vertex)
-> ST s (Vector Vertex -> Vector Vertex -> Hld)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Vector Vertex -> ST s (Vector Vertex)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector Vertex
revIndex
    ST s (Vector Vertex -> Vector Vertex -> Hld)
-> ST s (Vector Vertex) -> ST s (Vector Vertex -> Hld)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Vector Vertex -> ST s (Vector Vertex)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector Vertex
depths
    ST s (Vector Vertex -> Hld) -> ST s (Vector Vertex) -> ST s Hld
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Vector Vertex -> ST s (Vector Vertex)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector Vertex
subtreeSize
  where
    !n :: Vertex
n = Csr w -> Vertex
forall w. Csr w -> Vertex
Gr.nCsr Csr w
tree
    !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vertex
2 Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
* (Csr w -> Vertex
forall w. Csr w -> Vertex
Gr.nCsr Csr w
tree Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1) Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Csr w -> Vertex
forall w. Csr w -> Vertex
Gr.mCsr Csr w
tree) String
"AtCoder.Extra.Hld.newAt: not a non-directed tree"

-- | \(O(\log n)\) Calculates the lowest common ancestor of \(u\) and \(v\).
--
-- @since 1.1.0.0
{-# INLINE lca #-}
lca :: (HasCallStack) => Hld -> Vertex -> Vertex -> Vertex
lca :: HasCallStack => Hld -> Vertex -> Vertex -> Vertex
lca Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} = Vertex -> Vertex -> Vertex
inner
  where
    inner :: Vertex -> Vertex -> Vertex
inner !Vertex
x !Vertex
y
      -- sort for easier processing
      -- TODO: @case compare ix iy@ would be easier for me to understand
      | Vertex
ix Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
> Vertex
iy = Vertex -> Vertex -> Vertex
inner Vertex
y Vertex
x
      -- @x@ and @y@ are in other segments:
      | Vertex
hx Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= Vertex
hy = Vertex -> Vertex -> Vertex
inner Vertex
x (Vertex -> Vertex) -> Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Vector Vertex
parentHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hy
      -- @x@ and @y@ are within the same segment:
      -- select the smaller one, which is closer to the root and that is the LCA.
      | Bool
otherwise = Vertex
x
      where
        !ix :: Vertex
ix = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
x
        !iy :: Vertex
iy = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
y
        hx :: Vertex
hx = Vector Vertex
headHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
x
        hy :: Vertex
hy = Vector Vertex
headHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
y

-- | \(O(\log n)\) Go up \(k\) times from a vertex \(v\) to the root node. Throws an error if \(k\)
-- is bigger than the depth of \(v\).
--
-- @since 1.1.0.0
{-# INLINE ancestor #-}
ancestor :: (HasCallStack) => Hld -> Vertex -> Int -> Vertex
ancestor :: HasCallStack => Hld -> Vertex -> Vertex -> Vertex
ancestor Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
parent Vertex
k0 = Vertex -> Vertex -> Vertex
inner Vertex
parent Vertex
k0
  where
    !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vertex
0 Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
k0 Bool -> Bool -> Bool
&& Vertex
k0 Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
parent) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Tree.Hld.ancestor: k-th ancestor is out of the bounds (`k = " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Vertex -> String
forall a. Show a => a -> String
show Vertex
k0 String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"`)"
    inner :: Vertex -> Vertex -> Vertex
inner Vertex
v Vertex
k
      -- on this segment
      | Vertex
k Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
iv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
ihv = Vector Vertex
revIndexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! (Vertex
iv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
k)
      -- next segment
      | Bool
otherwise = Vertex -> Vertex -> Vertex
inner (Vector Vertex
parentHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hv) (Vertex
k Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- (Vertex
iv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
ihv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1))
      where
        iv :: Vertex
iv = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
v
        hv :: Vertex
hv = Vector Vertex
headHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
v
        ihv :: Vertex
ihv = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hv

-- | \(O(\log n)\) Returns the \(k\)-th vertex of the path between \(u\) and \(v\) from \(u\).
-- Throws an error if `k` is out
--
-- @since 1.1.0.0
{-# INLINE jump #-}
jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex
jump :: HasCallStack => Hld -> Vertex -> Vertex -> Vertex -> Maybe Vertex
jump hld :: Hld
hld@Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
u Vertex
v Vertex
k
  | Vertex
k Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
> Vertex
lenU Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
lenV = Maybe Vertex
forall a. Maybe a
Nothing
  | Vertex
k Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
lenU = Vertex -> Maybe Vertex
forall a. a -> Maybe a
Just (Vertex -> Maybe Vertex) -> Vertex -> Maybe Vertex
forall a b. (a -> b) -> a -> b
$ HasCallStack => Hld -> Vertex -> Vertex -> Vertex
Hld -> Vertex -> Vertex -> Vertex
ancestor Hld
hld Vertex
u Vertex
k
  | Bool
otherwise = Vertex -> Maybe Vertex
forall a. a -> Maybe a
Just (Vertex -> Maybe Vertex) -> Vertex -> Maybe Vertex
forall a b. (a -> b) -> a -> b
$ HasCallStack => Hld -> Vertex -> Vertex -> Vertex
Hld -> Vertex -> Vertex -> Vertex
ancestor Hld
hld Vertex
v (Vertex
lenU Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
lenV Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
k)
  where
    lca_ :: Vertex
lca_ = HasCallStack => Hld -> Vertex -> Vertex -> Vertex
Hld -> Vertex -> Vertex -> Vertex
lca Hld
hld Vertex
u Vertex
v
    du :: Vertex
du = Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
u
    dv :: Vertex
dv = Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
v
    lenU :: Vertex
lenU = Vertex
du Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
lca_
    lenV :: Vertex
lenV = Vertex
dv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
lca_

-- | \(O(\log n)\) Returns the length of the path between \(u\) and \(v\).
--
-- @since 1.1.0.0
{-# INLINE lengthBetween #-}
lengthBetween :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int
lengthBetween :: HasCallStack => Hld -> Vertex -> Vertex -> Vertex
lengthBetween hld :: Hld
hld@Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
u Vertex
v = Vertex
du Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
dLca Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
dv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
dLca
  where
    !lca_ :: Vertex
lca_ = HasCallStack => Hld -> Vertex -> Vertex -> Vertex
Hld -> Vertex -> Vertex -> Vertex
lca Hld
hld Vertex
u Vertex
v
    !dLca :: Vertex
dLca = Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
lca_
    !du :: Vertex
du = Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
u
    !dv :: Vertex
dv = Vector Vertex
depthHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
v

-- | \(O(n)\) Returns the vertices on the path between \(u\) and \(v\).
--
-- @since 1.1.0.0
{-# INLINE path #-}
path :: (HasCallStack) => Hld -> Vertex -> Vertex -> [Vertex]
path :: HasCallStack => Hld -> Vertex -> Vertex -> [Vertex]
path hld :: Hld
hld@Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
u Vertex
v = ((Vertex, Vertex) -> [Vertex]) -> [(Vertex, Vertex)] -> [Vertex]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Vertex, Vertex) -> [Vertex]
expand ([(Vertex, Vertex)] -> [Vertex]) -> [(Vertex, Vertex)] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ HasCallStack =>
WeightPolicy -> Hld -> Vertex -> Vertex -> [(Vertex, Vertex)]
WeightPolicy -> Hld -> Vertex -> Vertex -> [(Vertex, Vertex)]
pathSegmentsInclusive WeightPolicy
WeightsAreOnVertices Hld
hld Vertex
u Vertex
v
  where
    expand :: (Vertex, Vertex) -> [Vertex]
expand (!Vertex
l, !Vertex
r)
      | Vertex
l Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
r = (Vertex -> Vertex) -> [Vertex] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map (Vector Vertex
revIndexHld VG.!) [Vertex
l .. Vertex
r]
      | Bool
otherwise = (Vertex -> Vertex) -> [Vertex] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map (Vector Vertex
revIndexHld VG.!) [Vertex
l, Vertex
l Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1 .. Vertex
r]

-- | \(O(\log n)\) Decomposes a path between two vertices \(u\) and \(v\) into segments. Each
-- segment is represented as an __inclusive__ range \([u_i, v_i]\) of `VertexHLD`.
--
-- The LCA is omitted from the returning vertices when the weight policy is set to
-- `WeightsAreOnEdges`. This is the trick to put edge weights to on vertices.
--
-- @since 1.1.0.0
{-# INLINE pathSegmentsInclusive #-}
pathSegmentsInclusive :: (HasCallStack) => WeightPolicy -> Hld -> Vertex -> Vertex -> [(VertexHld, VertexHld)]
pathSegmentsInclusive :: HasCallStack =>
WeightPolicy -> Hld -> Vertex -> Vertex -> [(Vertex, Vertex)]
pathSegmentsInclusive WeightPolicy
weightPolicy Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
x0 Vertex
y0 = ([(Vertex, Vertex)], [(Vertex, Vertex)]) -> [(Vertex, Vertex)]
forall {a}. ([a], [a]) -> [a]
done (([(Vertex, Vertex)], [(Vertex, Vertex)]) -> [(Vertex, Vertex)])
-> ([(Vertex, Vertex)], [(Vertex, Vertex)]) -> [(Vertex, Vertex)]
forall a b. (a -> b) -> a -> b
$ Vertex
-> [(Vertex, Vertex)]
-> Vertex
-> [(Vertex, Vertex)]
-> ([(Vertex, Vertex)], [(Vertex, Vertex)])
inner Vertex
x0 [] Vertex
y0 []
  where
    isEdge :: Bool
isEdge = WeightPolicy
weightPolicy WeightPolicy -> WeightPolicy -> Bool
forall a. Eq a => a -> a -> Bool
== WeightPolicy
WeightsAreOnEdges
    done :: ([a], [a]) -> [a]
done (![a]
up, ![a]
down) = [a] -> [a]
forall a. [a] -> [a]
reverse [a]
up [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
down
    -- @up@: bottom to top. [(max, min)]
    -- @down@: top to bottom. [(min, max)]
    inner :: Vertex -> [(VertexHld, VertexHld)] -> Vertex -> [(VertexHld, VertexHld)] -> ([(VertexHld, VertexHld)], [(VertexHld, VertexHld)])
    inner :: Vertex
-> [(Vertex, Vertex)]
-> Vertex
-> [(Vertex, Vertex)]
-> ([(Vertex, Vertex)], [(Vertex, Vertex)])
inner Vertex
x [(Vertex, Vertex)]
up Vertex
y [(Vertex, Vertex)]
down
      | Vertex
hx Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
hy Bool -> Bool -> Bool
&& Bool
isEdge = case Vertex -> Vertex -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Vertex
ix Vertex
iy of
          -- skip LCA on edge vertices
          Ordering
LT -> ([(Vertex, Vertex)]
up, (Vertex
ix {- edge -} Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1, Vertex
iy) (Vertex, Vertex) -> [(Vertex, Vertex)] -> [(Vertex, Vertex)]
forall a. a -> [a] -> [a]
: [(Vertex, Vertex)]
down)
          Ordering
GT -> ((Vertex
ix, Vertex
iy {- edge -} Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) (Vertex, Vertex) -> [(Vertex, Vertex)] -> [(Vertex, Vertex)]
forall a. a -> [a] -> [a]
: [(Vertex, Vertex)]
up, [(Vertex, Vertex)]
down)
          Ordering
EQ -> ([(Vertex, Vertex)]
up, [(Vertex, Vertex)]
down)
      | Vertex
hx Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
hy Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
isEdge = case Vertex -> Vertex -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Vertex
ix Vertex
iy of
          Ordering
LT -> ([(Vertex, Vertex)]
up, (Vertex
ix, Vertex
iy) (Vertex, Vertex) -> [(Vertex, Vertex)] -> [(Vertex, Vertex)]
forall a. a -> [a] -> [a]
: [(Vertex, Vertex)]
down)
          Ordering
_ -> ((Vertex
ix, Vertex
iy) (Vertex, Vertex) -> [(Vertex, Vertex)] -> [(Vertex, Vertex)]
forall a. a -> [a] -> [a]
: [(Vertex, Vertex)]
up, [(Vertex, Vertex)]
down)
      | Bool
otherwise = case Vertex -> Vertex -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Vertex
ix Vertex
iy of
          Ordering
LT -> Vertex
-> [(Vertex, Vertex)]
-> Vertex
-> [(Vertex, Vertex)]
-> ([(Vertex, Vertex)], [(Vertex, Vertex)])
inner Vertex
x [(Vertex, Vertex)]
up Vertex
phy ((Vertex
ihy, Vertex
iy) (Vertex, Vertex) -> [(Vertex, Vertex)] -> [(Vertex, Vertex)]
forall a. a -> [a] -> [a]
: [(Vertex, Vertex)]
down)
          Ordering
GT -> Vertex
-> [(Vertex, Vertex)]
-> Vertex
-> [(Vertex, Vertex)]
-> ([(Vertex, Vertex)], [(Vertex, Vertex)])
inner Vertex
phx ((Vertex
ix, Vertex
ihx) (Vertex, Vertex) -> [(Vertex, Vertex)] -> [(Vertex, Vertex)]
forall a. a -> [a] -> [a]
: [(Vertex, Vertex)]
up) Vertex
y [(Vertex, Vertex)]
down
          Ordering
EQ -> String -> ([(Vertex, Vertex)], [(Vertex, Vertex)])
forall a. HasCallStack => String -> a
error String
"unreachable"
      where
        ix, iy :: VertexHld
        !ix :: Vertex
ix = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
x
        !iy :: Vertex
iy = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
y
        hx, hy :: Vertex
        hx :: Vertex
hx = Vector Vertex
headHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
x
        hy :: Vertex
hy = Vector Vertex
headHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
y
        ihx, ihy :: VertexHld
        ihx :: Vertex
ihx = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hx
        ihy :: Vertex
ihy = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hy
        phx, phy :: VertexHld
        phx :: Vertex
phx = Vector Vertex
parentHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hx
        phy :: Vertex
phy = Vector Vertex
parentHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
hy

-- | \(O(1)\) Returns a half-open interval of `VertexHld` \([\mathrm{start}, \mathrm{end})\) that
-- corresponds to the subtree segments rooted at the given @subtreeRoot@.
--
-- @since 1.1.0.0
{-# INLINE subtreeSegmentInclusive #-}
subtreeSegmentInclusive :: (HasCallStack) => Hld -> Vertex -> (VertexHld, VertexHld)
subtreeSegmentInclusive :: HasCallStack => Hld -> Vertex -> (Vertex, Vertex)
subtreeSegmentInclusive Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
subtreeRoot = (Vertex
ir, Vertex
ir Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
sr Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1)
  where
    ir :: Vertex
ir = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
subtreeRoot
    sr :: Vertex
sr = Vector Vertex
subtreeSizeHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
subtreeRoot

-- | \(O(1)\) Returns `True` if \(u\) is in a subtree of \(r\).
--
-- @since 1.1.0.0
{-# INLINE isInSubtree #-}
isInSubtree :: (HasCallStack) => Hld -> Vertex -> Vertex -> Bool
isInSubtree :: HasCallStack => Hld -> Vertex -> Vertex -> Bool
isInSubtree hld :: Hld
hld@Hld {Vertex
Vector Vertex
indexHld :: Hld -> Vector Vertex
headHld :: Hld -> Vector Vertex
parentHld :: Hld -> Vector Vertex
rootHld :: Hld -> Vertex
revIndexHld :: Hld -> Vector Vertex
depthHld :: Hld -> Vector Vertex
subtreeSizeHld :: Hld -> Vector Vertex
rootHld :: Vertex
parentHld :: Vector Vertex
indexHld :: Vector Vertex
headHld :: Vector Vertex
revIndexHld :: Vector Vertex
depthHld :: Vector Vertex
subtreeSizeHld :: Vector Vertex
..} Vertex
r_ Vertex
u = Vertex
l Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
iu Bool -> Bool -> Bool
&& Vertex
iu Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
r
  where
    (!Vertex
l, !Vertex
r) = HasCallStack => Hld -> Vertex -> (Vertex, Vertex)
Hld -> Vertex -> (Vertex, Vertex)
subtreeSegmentInclusive Hld
hld Vertex
r_
    !iu :: Vertex
iu = Vector Vertex
indexHld Vector Vertex -> Vertex -> Vertex
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Vertex -> a
VG.! Vertex
u

-- | Represents whether weights are put on vertices or edges.
--
-- @since 1.1.0.0
data WeightPolicy
  = -- | Weights are put on vertices.
    --
    -- @since 1.1.0.0
    WeightsAreOnVertices
  | -- | Weights are put on edges.
    --
    -- @since 1.1.0.0
    WeightsAreOnEdges
  deriving
    ( -- | @since 1.1.0.0
      WeightPolicy -> WeightPolicy -> Bool
(WeightPolicy -> WeightPolicy -> Bool)
-> (WeightPolicy -> WeightPolicy -> Bool) -> Eq WeightPolicy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WeightPolicy -> WeightPolicy -> Bool
== :: WeightPolicy -> WeightPolicy -> Bool
$c/= :: WeightPolicy -> WeightPolicy -> Bool
/= :: WeightPolicy -> WeightPolicy -> Bool
Eq,
      -- | @since 1.1.0.0
      Vertex -> WeightPolicy -> ShowS
[WeightPolicy] -> ShowS
WeightPolicy -> String
(Vertex -> WeightPolicy -> ShowS)
-> (WeightPolicy -> String)
-> ([WeightPolicy] -> ShowS)
-> Show WeightPolicy
forall a.
(Vertex -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Vertex -> WeightPolicy -> ShowS
showsPrec :: Vertex -> WeightPolicy -> ShowS
$cshow :: WeightPolicy -> String
show :: WeightPolicy -> String
$cshowList :: [WeightPolicy] -> ShowS
showList :: [WeightPolicy] -> ShowS
Show
    )

-- | \(O(\log n f)\) Returns product of the path between \(u\) and \(v\), using the user functions
-- of time complexity \(O(f)\).
--
-- @since 1.1.0.0
{-# INLINE prod #-}
prod ::
  (HasCallStack, Monoid mono, Monad m) =>
  -- | The `WeightPolicy`.
  WeightPolicy ->
  -- | The `Hld`.
  Hld ->
  -- | User function for getting products in \([u, v)\), where \(u < v\) and
  -- \(\mathrm{depth}(u) < \mathrm{depth}(v)\).
  (VertexHld -> VertexHld -> m mono) ->
  -- | User function for getting products in \([u, v)\), where \(u < v\) and
  -- \(\mathrm{depth}(u) > \mathrm{depth}(v)\).
  (VertexHld -> VertexHld -> m mono) ->
  -- | \(u\).
  Vertex ->
  -- | \(v\).
  Vertex ->
  -- | Product of the path between \(u\) and \(v\).
  m mono
prod :: forall mono (m :: * -> *).
(HasCallStack, Monoid mono, Monad m) =>
WeightPolicy
-> Hld
-> (Vertex -> Vertex -> m mono)
-> (Vertex -> Vertex -> m mono)
-> Vertex
-> Vertex
-> m mono
prod WeightPolicy
weightPolicy Hld
hld Vertex -> Vertex -> m mono
prodF Vertex -> Vertex -> m mono
prodB Vertex
u0 Vertex
v0 = do
  (mono -> (Vertex, Vertex) -> m mono)
-> mono -> [(Vertex, Vertex)] -> m mono
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
    ( \ !mono
acc (!Vertex
u, !Vertex
v) -> do
        !mono
x <-
          if Vertex
u Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
v
            then Vertex -> Vertex -> m mono
prodF Vertex
u (Vertex -> m mono) -> Vertex -> m mono
forall a b. (a -> b) -> a -> b
$ Vertex
v Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1
            else Vertex -> Vertex -> m mono
prodB Vertex
v (Vertex -> m mono) -> Vertex -> m mono
forall a b. (a -> b) -> a -> b
$ Vertex
u Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1
        mono -> m mono
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (mono -> m mono) -> mono -> m mono
forall a b. (a -> b) -> a -> b
$! mono
acc mono -> mono -> mono
forall a. Semigroup a => a -> a -> a
<> mono
x
    )
    mono
forall a. Monoid a => a
mempty
    (HasCallStack =>
WeightPolicy -> Hld -> Vertex -> Vertex -> [(Vertex, Vertex)]
WeightPolicy -> Hld -> Vertex -> Vertex -> [(Vertex, Vertex)]
pathSegmentsInclusive WeightPolicy
weightPolicy Hld
hld Vertex
u0 Vertex
v0)