{-# LANGUAGE RecordWildCards #-}

-- | Integration of segment trees with the heavy-light decomposition technique. Computes monoid
-- products on a path in \(O(\log^2 n)\) time or on a subtree in \(O(\log n)\) time.
--
-- - If vertices have weights, create a `TreeMonoid` with `fromVerts`.
-- - If edges have weights, create a tree monoid with `fromEdges`.
--
-- ==== __(Internals) Weights on edges__
--
-- When vertices are unweighted and only edges have weights, treat edges as new vertices or assign
-- edge weights to the deeper vertex.
--
-- Idea 1. Convert edges into new vertices. This is inefficient.
--
-- @
-- o--o--o  --> o-x-o-x-o
-- @
--
-- Idea 2. Assign edge weight to the deeper vertex. The is the internal implementation of
-- `fromEdges` and LCAs are ignored on `prod`:
--
-- @
--   o
--   | <--- edge 1
--   o <- write weight 1 here
--   | <--- edge 2
--   o <- write weight 2 here
-- @
--
-- ==== __Example (1): Weights are on vertices__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld
-- >>> import AtCoder.Extra.Tree.TreeMonoid qualified as TM
-- >>> import Data.Semigroup (Sum (..))
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0--1--2--3
-- >>> --    +
-- >>> --    +--4--5
-- >>> let n = 6
-- >>> -- note that the edges must be bi-directed:
-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)]
-- >>> let weights = VU.generate n Sum -- vertex `i` is given weight of `i`
-- >>> let hld = Hld.new tree
-- >>> tm <- TM.fromVerts hld {- `Sum` is commutative -} Commute weights
-- >>> TM.prod tm 1 3
-- Sum {getSum = 6}
--
-- >>> TM.prodSubtree tm 1
-- Sum {getSum = 15}
--
-- >>> TM.write tm 1 $ Sum 10
-- >>> TM.prod tm 1 3
-- Sum {getSum = 15}
--
-- ==== __Example (2): Weights are on edges__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld
-- >>> import AtCoder.Extra.Tree.TreeMonoid qualified as TM
-- >>> import Data.Semigroup (Sum (..))
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0--1--2--3
-- >>> --    +
-- >>> --    +--4--5
-- >>> let n = 6
-- >>> let edges = VU.fromList [(0, 1, Sum (1 :: Int)), (1, 2, Sum 2), (2, 3, Sum 3), (1, 4, Sum 4), (4, 5, Sum 5)]
-- >>> -- note that the edges must be bi-directed:
-- >>> let tree = Gr.build n $ Gr.swapDupe edges
-- >>> let hld = Hld.new tree
-- >>> let hld = Hld.new tree
-- >>> -- note that the edge doesn't have to be bi-directed:
-- >>> tm <- TM.fromEdges hld {- `Sum` is commutative -} Commute edges
-- >>> TM.prod tm 1 3
-- Sum {getSum = 5}
--
-- >>> TM.prodSubtree tm 1
-- Sum {getSum = 14}
--
-- >>> TM.write tm 2 $ Sum 10
-- >>> TM.prod tm 1 3
-- Sum {getSum = 13}
--
-- @since 1.1.0.0
module AtCoder.Extra.Tree.TreeMonoid
  ( -- * TreeMonoid
    TreeMonoid,
    Vertex,
    VertexHld,
    Commutativity (..),

    -- * Constructors
    fromVerts,
    fromEdges,

    -- * Segment tree methods

    -- ** Reading
    prod,
    prodSubtree,
    read,

    -- ** Modifications
    write,
    exchange,
    modify,
    modifyM,
  )
where

import AtCoder.Extra.Tree.Hld qualified as Hld
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.SegTree qualified as ST
import Control.Monad
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Monoid (Dual (..))
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)
import Prelude hiding (read)

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

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

-- | A wrapper for `Hld` getting product on paths on a tree using `Hld` and segment tree(s).
--
-- @since 1.1.0.0
data TreeMonoid a s = TreeMonoid
  { -- | Borrowed Hld.
    forall a s. TreeMonoid a s -> Hld
hldTM :: !Hld.Hld,
    -- | Indicates if it's targetting commutative monoids.
    forall a s. TreeMonoid a s -> Commutativity
commuteTM :: !Commutativity,
    -- | Indicates if it's targetting edge weights (If not, it's targetting vertex weights).
    forall a s. TreeMonoid a s -> WeightPolicy
weightPolicyTM :: !Hld.WeightPolicy,
    -- | Segment tree for getting products upwards.
    forall a s. TreeMonoid a s -> SegTree s a
segFTM :: !(ST.SegTree s a),
    -- | Segment tree for getting products downwards. Only created when the monoid is
    -- `NonCommute`.
    forall a s. TreeMonoid a s -> SegTree s (Dual a)
segBTM :: !(ST.SegTree s (Dual a))
  }

-- | Represents whether a monoid is commutative or noncommutative.
--
-- @since 1.1.0.0
data Commutativity
  = -- | Commutative: \(a \cdot b = b \cdot a\).
    --
    -- @since 1.1.0.0
    Commute
  | -- | Noncommutative: \(a \cdot b \neq b \cdot a\).
    --
    -- @since 1.1.0.0
    NonCommute
  deriving
    ( -- | @since 1.1.0.0
      Commutativity -> Commutativity -> Bool
(Commutativity -> Commutativity -> Bool)
-> (Commutativity -> Commutativity -> Bool) -> Eq Commutativity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Commutativity -> Commutativity -> Bool
== :: Commutativity -> Commutativity -> Bool
$c/= :: Commutativity -> Commutativity -> Bool
/= :: Commutativity -> Commutativity -> Bool
Eq,
      -- | @since 1.1.0.0
      Int -> Commutativity -> ShowS
[Commutativity] -> ShowS
Commutativity -> String
(Int -> Commutativity -> ShowS)
-> (Commutativity -> String)
-> ([Commutativity] -> ShowS)
-> Show Commutativity
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Commutativity -> ShowS
showsPrec :: Int -> Commutativity -> ShowS
$cshow :: Commutativity -> String
show :: Commutativity -> String
$cshowList :: [Commutativity] -> ShowS
showList :: [Commutativity] -> ShowS
Show
    )

-- | \(O(n)\) Creates a `TreeMonoid` with weights on vertices.
--
-- @since 1.1.0.0
{-# INLINE fromVerts #-}
fromVerts ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | `Hld.Hld`.
  Hld.Hld ->
  -- | Whether the monoid is commutative or not.
  Commutativity ->
  -- | The vertex weights.
  VU.Vector a ->
  -- | A `TreeMonoid` with weights on vertices.
  m (TreeMonoid a (PrimState m))
fromVerts :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld -> Commutativity -> Vector a -> m (TreeMonoid a (PrimState m))
fromVerts Hld
hld Commutativity
commuteTM Vector a
xs_ = ST (PrimState m) (TreeMonoid a (PrimState m))
-> m (TreeMonoid a (PrimState m))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (TreeMonoid a (PrimState m))
 -> m (TreeMonoid a (PrimState m)))
-> ST (PrimState m) (TreeMonoid a (PrimState m))
-> m (TreeMonoid a (PrimState m))
forall a b. (a -> b) -> a -> b
$ Hld
-> Commutativity
-> Vector a
-> ST (PrimState m) (TreeMonoid a (PrimState m))
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld -> Commutativity -> Vector a -> ST s (TreeMonoid a s)
fromVertsST Hld
hld Commutativity
commuteTM Vector a
xs_

-- | \(O(n)\) Creates a `TreeMonoid` with weignts on edges. The edges are not required to be
-- duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is needed.
--
-- @since 1.1.0.0
{-# INLINE fromEdges #-}
fromEdges ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | `Hld.Hld`.
  Hld.Hld ->
  -- | Whether the monoid is commutative or not.
  Commutativity ->
  -- | Input edges.
  VU.Vector (Vertex, Vertex, a) ->
  -- | A `TreeMonoid` with weights on edges.
  m (TreeMonoid a (PrimState m))
fromEdges :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> Vector (Int, Int, a)
-> m (TreeMonoid a (PrimState m))
fromEdges Hld
hld Commutativity
commuteTM Vector (Int, Int, a)
edges = ST (PrimState m) (TreeMonoid a (PrimState m))
-> m (TreeMonoid a (PrimState m))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (TreeMonoid a (PrimState m))
 -> m (TreeMonoid a (PrimState m)))
-> ST (PrimState m) (TreeMonoid a (PrimState m))
-> m (TreeMonoid a (PrimState m))
forall a b. (a -> b) -> a -> b
$ Hld
-> Commutativity
-> Vector (Int, Int, a)
-> ST (PrimState m) (TreeMonoid a (PrimState m))
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld
-> Commutativity -> Vector (Int, Int, a) -> ST s (TreeMonoid a s)
fromEdgesST Hld
hld Commutativity
commuteTM Vector (Int, Int, a)
edges

-- | \(O(\log^2 n)\) Returns the product of the path between two vertices \(u\), \(v\) (invlusive).
--
-- @since 1.1.0.0
{-# INLINE prod #-}
prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> Vertex -> m a
prod :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> Int -> m a
prod TreeMonoid a (PrimState m)
tm Int
u Int
v = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ TreeMonoid a (PrimState m) -> Int -> Int -> ST (PrimState m) a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> Int -> ST s a
prodST TreeMonoid a (PrimState m)
tm Int
u Int
v

-- | \(O(\log n)\) Returns the product of the subtree rooted at the given `Vertex`.
--
-- @since 1.1.0.0
{-# INLINE prodSubtree #-}
prodSubtree :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a
prodSubtree :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> m a
prodSubtree TreeMonoid a (PrimState m)
tm Int
subtreeRoot = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ TreeMonoid a (PrimState m) -> Int -> ST (PrimState m) a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> ST s a
prodSubtreeST TreeMonoid a (PrimState m)
tm Int
subtreeRoot

-- | \(O(1)\) Reads a monoid value of a `Vertex`.
--
-- @since 1.1.0.0
{-# INLINE read #-}
read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a
read :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> m a
read TreeMonoid a (PrimState m)
tm Int
i_ = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ TreeMonoid a (PrimState m) -> Int -> ST (PrimState m) a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> ST s a
readST TreeMonoid a (PrimState m)
tm Int
i_

-- | \(O(\log n)\) Writes to the monoid value of a vertex.
--
-- @since 1.1.0.0
{-# INLINE write #-}
write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> a -> m ()
write TreeMonoid a (PrimState m)
tm Int
i_ a
x = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ TreeMonoid a (PrimState m) -> Int -> a -> ST (PrimState m) ()
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> a -> ST s ()
writeST TreeMonoid a (PrimState m)
tm Int
i_ a
x

-- | \(O(\log n)\) Writes to the monoid value of a vertex and returns the old value.
--
-- @since 1.1.0.0
{-# INLINE exchange #-}
exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m a
exchange :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> a -> m a
exchange TreeMonoid a (PrimState m)
tm Int
i_ a
x = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ TreeMonoid a (PrimState m) -> Int -> a -> ST (PrimState m) a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> a -> ST s a
exchangeST TreeMonoid a (PrimState m)
tm Int
i_ a
x

-- | \(O(\log n)\) Modifies the monoid value of a vertex with a pure function.
--
-- @since 1.1.0.0
{-# INLINE modify #-}
modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m ()
modify TreeMonoid a (PrimState m)
tm a -> a
f Int
i_ = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ TreeMonoid a (PrimState m)
-> (a -> a) -> Int -> ST (PrimState m) ()
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> (a -> a) -> Int -> ST s ()
modifyST TreeMonoid a (PrimState m)
tm a -> a
f Int
i_

-- | \(O(\log n)\) Modifies the monoid value of a vertex with a monadic function.
--
-- @since 1.1.0.0
{-# INLINE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> m a) -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> (a -> m a) -> Int -> m ()
modifyM TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} a -> m a
f Int
i_ = do
  let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
  SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
ST.modifyM SegTree (PrimState m) a
segFTM a -> m a
f Int
i
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    SegTree (PrimState m) (Dual a)
-> (Dual a -> m (Dual a)) -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
ST.modifyM SegTree (PrimState m) (Dual a)
segBTM ((a -> Dual a
forall a. a -> Dual a
Dual <$>) (m a -> m (Dual a)) -> (Dual a -> m a) -> Dual a -> m (Dual a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m a
f (a -> m a) -> (Dual a -> a) -> Dual a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual a -> a
forall a. Dual a -> a
getDual) Int
i

-- -------------------------------------------------------------------------------------------------
-- INLINE
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE buildST #-}
buildST ::
  (HasCallStack, Monoid a, VU.Unbox a) =>
  Hld.Hld ->
  Commutativity ->
  Hld.WeightPolicy ->
  VU.Vector a ->
  ST s (TreeMonoid a s)
buildST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> ST s (TreeMonoid a s)
buildST Hld
hldTM Commutativity
commuteTM WeightPolicy
weightPolicyTM Vector a
weights = do
  SegTree s a
segFTM <- Vector a -> ST s (SegTree (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
ST.build Vector a
weights
  SegTree s (Dual a)
segBTM <-
    case Commutativity
commuteTM of
      Commutativity
Commute -> Vector (Dual a) -> ST s (SegTree (PrimState (ST s)) (Dual a))
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
ST.build Vector (Dual a)
forall a. Unbox a => Vector a
VU.empty
      Commutativity
NonCommute -> Vector (Dual a) -> ST s (SegTree (PrimState (ST s)) (Dual a))
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
ST.build (Vector (Dual a) -> ST s (SegTree (PrimState (ST s)) (Dual a)))
-> Vector (Dual a) -> ST s (SegTree (PrimState (ST s)) (Dual a))
forall a b. (a -> b) -> a -> b
$ (a -> Dual a) -> Vector a -> Vector (Dual a)
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map a -> Dual a
forall a. a -> Dual a
Dual Vector a
weights
  TreeMonoid a s -> ST s (TreeMonoid a s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..}

{-# INLINEABLE fromVertsST #-}
fromVertsST ::
  (HasCallStack, Monoid a, VU.Unbox a) =>
  Hld.Hld ->
  Commutativity ->
  VU.Vector a ->
  ST s (TreeMonoid a s)
fromVertsST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld -> Commutativity -> Vector a -> ST s (TreeMonoid a s)
fromVertsST hld :: Hld
hld@Hld.Hld {Vector Int
indexHld :: Hld -> Vector Int
indexHld :: Vector Int
indexHld} Commutativity
commuteTM Vector a
xs_ = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexHld Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs_) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Tree.TreeMonoid.fromVertsST: vertex number mismatch (`" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexHld) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"` and `" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs_) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"`)"
  let !xs :: Vector a
xs = (forall s. ST s (MVector s a)) -> Vector a
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s a)) -> Vector a)
-> (forall s. ST s (MVector s a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
        MVector s a
vec <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> ST s (MVector (PrimState (ST s)) a))
-> Int -> ST s (MVector (PrimState (ST s)) a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs_
        Vector a -> (Int -> a -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector a
xs_ ((Int -> a -> ST s ()) -> ST s ())
-> (Int -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i a
x -> do
          MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
vec (Vector Int
indexHld Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i) a
x
        MVector s a -> ST s (MVector s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s a
vec
  Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> ST s (TreeMonoid a s)
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> ST s (TreeMonoid a s)
buildST Hld
hld Commutativity
commuteTM WeightPolicy
Hld.WeightsAreOnVertices Vector a
xs

{-# INLINEABLE fromEdgesST #-}
fromEdgesST ::
  (HasCallStack, Monoid a, VU.Unbox a) =>
  Hld.Hld ->
  Commutativity ->
  VU.Vector (Vertex, Vertex, a) ->
  ST s (TreeMonoid a s)
fromEdgesST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld
-> Commutativity -> Vector (Int, Int, a) -> ST s (TreeMonoid a s)
fromEdgesST hld :: Hld
hld@Hld.Hld {Vector Int
indexHld :: Hld -> Vector Int
indexHld :: Vector Int
indexHld} Commutativity
commuteTM Vector (Int, Int, a)
edges = do
  let !xs :: Vector a
xs = (forall s. ST s (MVector s a)) -> Vector a
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s a)) -> Vector a)
-> (forall s. ST s (MVector s a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
        MVector s a
vec <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> ST s (MVector (PrimState (ST s)) a))
-> Int -> ST s (MVector (PrimState (ST s)) a)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexHld
        Vector (Int, Int, a) -> ((Int, Int, a) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, Int, a)
edges (((Int, Int, a) -> ST s ()) -> ST s ())
-> ((Int, Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
u, !Int
v, !a
w) -> do
          let u' :: Int
u' = Vector Int
indexHld Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
u
          let v' :: Int
v' = Vector Int
indexHld Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v
          MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
vec (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
u' Int
v') a
w
        MVector s a -> ST s (MVector s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s a
vec
  Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> ST s (TreeMonoid a s)
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> ST s (TreeMonoid a s)
buildST Hld
hld Commutativity
commuteTM WeightPolicy
Hld.WeightsAreOnEdges Vector a
xs

{-# INLINEABLE prodST #-}
prodST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> Vertex -> ST s a
prodST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> Int -> ST s a
prodST TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..} Int
u Int
v = do
  case Commutativity
commuteTM of
    Commutativity
Commute -> WeightPolicy
-> Hld
-> (Int -> Int -> ST s a)
-> (Int -> Int -> ST s a)
-> Int
-> Int
-> ST s a
forall mono (m :: * -> *).
(HasCallStack, Monoid mono, Monad m) =>
WeightPolicy
-> Hld
-> (Int -> Int -> m mono)
-> (Int -> Int -> m mono)
-> Int
-> Int
-> m mono
Hld.prod WeightPolicy
weightPolicyTM Hld
hldTM (SegTree (PrimState (ST s)) a -> Int -> Int -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree s a
SegTree (PrimState (ST s)) a
segFTM) (SegTree (PrimState (ST s)) a -> Int -> Int -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree s a
SegTree (PrimState (ST s)) a
segFTM) Int
u Int
v
    Commutativity
NonCommute -> WeightPolicy
-> Hld
-> (Int -> Int -> ST s a)
-> (Int -> Int -> ST s a)
-> Int
-> Int
-> ST s a
forall mono (m :: * -> *).
(HasCallStack, Monoid mono, Monad m) =>
WeightPolicy
-> Hld
-> (Int -> Int -> m mono)
-> (Int -> Int -> m mono)
-> Int
-> Int
-> m mono
Hld.prod WeightPolicy
weightPolicyTM Hld
hldTM (SegTree (PrimState (ST s)) a -> Int -> Int -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree s a
SegTree (PrimState (ST s)) a
segFTM) (\Int
l Int
r -> Dual a -> a
forall a. Dual a -> a
getDual (Dual a -> a) -> ST s (Dual a) -> ST s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegTree (PrimState (ST s)) (Dual a) -> Int -> Int -> ST s (Dual a)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree s (Dual a)
SegTree (PrimState (ST s)) (Dual a)
segBTM Int
l Int
r) Int
u Int
v

{-# INLINEABLE prodSubtreeST #-}
prodSubtreeST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> ST s a
prodSubtreeST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> ST s a
prodSubtreeST TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..} Int
subtreeRoot = do
  let (!Int
l, !Int
r) = HasCallStack => Hld -> Int -> (Int, Int)
Hld -> Int -> (Int, Int)
Hld.subtreeSegmentInclusive Hld
hldTM Int
subtreeRoot
  case WeightPolicy
weightPolicyTM of
    WeightPolicy
Hld.WeightsAreOnVertices -> SegTree (PrimState (ST s)) a -> Int -> Int -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree s a
SegTree (PrimState (ST s)) a
segFTM Int
l (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    WeightPolicy
Hld.WeightsAreOnEdges -> do
      -- ignore the root of the subtree, which has the minimum index among the subtree vertices
      if Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r
        then a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
        else SegTree (PrimState (ST s)) a -> Int -> Int -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree s a
SegTree (PrimState (ST s)) a
segFTM (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINEABLE readST #-}
readST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> ST s a
readST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> ST s a
readST TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..} Int
i_ = do
  let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
  SegTree (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m a
ST.read SegTree s a
SegTree (PrimState (ST s)) a
segFTM Int
i

{-# INLINEABLE writeST #-}
writeST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> a -> ST s ()
writeST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> a -> ST s ()
writeST TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..} Int
i_ a
x = do
  let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
  SegTree (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
ST.write SegTree s a
SegTree (PrimState (ST s)) a
segFTM Int
i a
x
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    SegTree (PrimState (ST s)) (Dual a) -> Int -> Dual a -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
ST.write SegTree s (Dual a)
SegTree (PrimState (ST s)) (Dual a)
segBTM Int
i (Dual a -> ST s ()) -> Dual a -> ST s ()
forall a b. (a -> b) -> a -> b
$ a -> Dual a
forall a. a -> Dual a
Dual a
x

{-# INLINEABLE exchangeST #-}
exchangeST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> a -> ST s a
exchangeST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> Int -> a -> ST s a
exchangeST TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..} Int
i_ a
x = do
  let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
  !a
res <- SegTree (PrimState (ST s)) a -> Int -> a -> ST s a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m a
ST.exchange SegTree s a
SegTree (PrimState (ST s)) a
segFTM Int
i a
x
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    SegTree (PrimState (ST s)) (Dual a) -> Int -> Dual a -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
ST.write SegTree s (Dual a)
SegTree (PrimState (ST s)) (Dual a)
segBTM Int
i (Dual a -> ST s ()) -> Dual a -> ST s ()
forall a b. (a -> b) -> a -> b
$ a -> Dual a
forall a. a -> Dual a
Dual a
x
  a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res

{-# INLINEABLE modifyST #-}
modifyST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> (a -> a) -> Int -> ST s ()
modifyST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
TreeMonoid a s -> (a -> a) -> Int -> ST s ()
modifyST TreeMonoid {WeightPolicy
Hld
SegTree s a
SegTree s (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree s a
segBTM :: SegTree s (Dual a)
..} a -> a
f Int
i_ = do
  let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
  SegTree (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
ST.modify SegTree s a
SegTree (PrimState (ST s)) a
segFTM a -> a
f Int
i
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    SegTree (PrimState (ST s)) (Dual a)
-> (Dual a -> Dual a) -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
ST.modify SegTree s (Dual a)
SegTree (PrimState (ST s)) (Dual a)
segBTM (a -> Dual a
forall a. a -> Dual a
Dual (a -> Dual a) -> (Dual a -> a) -> Dual a -> Dual a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f (a -> a) -> (Dual a -> a) -> Dual a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual a -> a
forall a. Dual a -> a
getDual) Int
i