content: -- | Generic tree functions.
  --
  -- @since 1.1.0.0
  module AtCoder.Extra.Tree
    ( -- * Tree properties
      diameter,
      diameterPath,
  
      -- * Minimum spanning tree
      mst,
      mstBy,
  
      -- * Tree folding
  
      -- | These functions 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,
      scan,
      foldReroot,
    )
  where
  
  import AtCoder.Dsu qualified as Dsu
  import AtCoder.Extra.Graph qualified as Gr
  import AtCoder.Extra.Vector.Prim qualified as EVP
  import Control.Monad (when)
  import Control.Monad.ST (runST)
  import Data.Bit (Bit (..))
  import Data.Functor.Identity (runIdentity)
  import Data.Maybe (isJust)
  import Data.Ord (comparing)
  import Data.Vector.Algorithms.Intro qualified as VAI
  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)
  
  -- | \(O(n + m)\) Returns the endpoints of the diameter of a tree and their distance: \(((u, v), w)\).
  --
  -- ==== __Example__
  -- >>> 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
  {-# INLINEABLE diameter #-}
  diameter ::
    (HasCallStack, VU.Unbox w, Num w, Ord w) =>
    -- | The number of vertices.
    Int ->
    -- | Graph given as a function.
    (Int -> VU.Vector (Int, w)) ->
    -- | Distances assigned to unreachable vertices.
    w ->
    -- | Tuple of (endpoints of the longest path in a tree, distance of it).
    ((Int, Int), w)
  diameter n gr !undefW =
    let !bfs1 = Gr.bfs n gr undefW $ VU.singleton (0, 0)
        !from = VU.maxIndex bfs1
        !bfs2 = Gr.bfs n gr undefW $ VU.singleton (from, 0)
        !to = VU.maxIndex bfs2
        !w = VU.maximum bfs2
     in ((from, to), w)
  
  -- | \(O(n + m)\) Returns the longest path in a tree and the distance of it.
  --
  -- ==== __Example__
  -- >>> 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
  {-# INLINEABLE diameterPath #-}
  diameterPath ::
    (HasCallStack, Show w, VU.Unbox w, Num w, Ord w) =>
    -- | The number of vertices.
    Int ->
    -- | Graph given as a function.
    (Int -> VU.Vector (Int, w)) ->
    -- | Distances assigned to unreachable vertices.
    w ->
    -- | Tuple of (the longest path, distance of it).
    (VU.Vector Int, w)
  diameterPath n gr !undefW =
    let !bfs1 = Gr.bfs n gr undefW $ VU.singleton (0, 0)
        !from = VU.maxIndex bfs1
        (!bfs2, !parents) = Gr.trackingBfs n gr undefW $ VU.singleton (from, 0)
        !to = VU.maxIndex bfs2
        !w = bfs2 VG.! to
     in (Gr.constructPathFromRoot parents to, w)
  
  -- | \(O(m \log m)\) Kruskal's algorithm. Returns \((w, \mathrm{use}, \mathrm{mst})\): (weight of the minimum
  -- spanning tree, edge use flags, minimum spanning tree).
  --
  -- NOTE: The edges need not be duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is required
  -- for each edge.
  --
  -- ==== __Example__
  -- 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
  {-# INLINE mst #-}
  mst :: (Num w, Ord w, VU.Unbox w) => Int -> VU.Vector (Int, Int, w) -> (w, VU.Vector Bit, Gr.Csr w)
  mst = mstBy (comparing id)
  
  -- | \(O(m \log m)\) Kruskal's algorithm for creating either minimum or maximum spanning tree. Returns
  -- \((w, \mathrm{use}, \mathrm{mst})\): (weight of the minimum spanning tree, edge use flags, minimum
  -- spanning tree).
  --
  --
  -- NOTE: The edges need not be duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is required
  -- for each edge.
  --
  -- ==== __Example__
  -- 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
  {-# INLINEABLE mstBy #-}
  mstBy :: (Num w, Ord w, VU.Unbox w) => (w -> w -> Ordering) -> Int -> VU.Vector (Int, Int, w) -> (w, VU.Vector Bit, Gr.Csr w)
  mstBy !f nVerts edges = runST $ do
    dsu <- Dsu.new nVerts
    wSum <- VUM.replicate 1 0
    use <-
      ( VU.accumulate
          (const id)
          (VU.replicate (VU.length edges) (Bit False))
          <$>
        )
        . EVP.mapM
          ( \(i :: Int) -> do
              let !u = us VG.! i
              let !v = vs VG.! i
              b <- isJust <$> Dsu.mergeMaybe dsu u v
              when b $ do
                VGM.modify wSum (+ ws VG.! i) 0
              pure (i, Bit b)
          )
        . VU.modify (VAI.sortBy (\(i :: Int) (j :: Int) -> f (ws VG.! i) (ws VG.! j)))
        $ VU.generate (VU.length edges) id
    let !gr = Gr.build nVerts $ Gr.swapDupe $ VU.ifilter (\i _ -> unBit (use VG.! i)) edges
    (,use,gr) <$> VGM.read wSum 0
    where
      (!us, !vs, !ws) = VU.unzip3 edges
  
  {-# INLINEABLE foldImpl #-}
  foldImpl ::
    forall m w f a.
    (HasCallStack, Monad m, VU.Unbox w) =>
    (Int -> VU.Vector (Int, w)) ->
    (Int -> a) ->
    (a -> (Int, w) -> f) ->
    (f -> a -> a) ->
    Int ->
    (Int -> a -> m ()) ->
    m a
  foldImpl tree valAt toF act root memo = inner (-1) root
    where
      inner :: Int -> Int -> m a
      inner !parent !v1 = do
        let !acc0 = valAt v1
        let !v2s = VU.filter ((/= parent) . fst) $ tree v1
        !res <- VU.foldM' (\acc (!v2, !w) -> (`act` acc) . (`toF` (v1, w)) <$> inner v1 v2) acc0 v2s
        memo v1 res
        pure res
  
  -- | \(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
  {-# INLINE fold #-}
  fold ::
    (HasCallStack, VU.Unbox w) =>
    -- | Graph as a function.
    (Int -> VU.Vector (Int, w)) ->
    -- | @valAt@: Assignment of initial vertex values.
    (Int -> a) ->
    -- | @toF@: Converts a vertex value into an action onto a neighbor vertex.
    (a -> (Int, w) -> f) ->
    -- | @act@: Performs an action onto a vertex value.
    (f -> a -> a) ->
    -- | Root vertex.
    Int ->
    -- | Tree folding result from the root vertex.
    a
  fold tree valAt toF act root = runIdentity $ do
    foldImpl tree valAt toF act root (\_ _ -> pure ())
  
  -- | \(O(n)\) Folds a tree from a root vertex, also known as tree DP. The calculation process on
  -- every vertex is recorded 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.3.0.0
  {-# INLINE scan #-}
  scan ::
    (VU.Unbox w, VU.Unbox a) =>
    -- | The number of vertices.
    Int ->
    -- | Graph as a function.
    (Int -> VU.Vector (Int, w)) ->
    -- | @valAt@: Assignment of initial vertex values.
    (Int -> a) ->
    -- | @toF@: Converts a vertex value into an action onto a neighbor vertex.
    (a -> (Int, w) -> f) ->
    -- | @act@: Performs an action onto a vertex value.
    (f -> a -> a) ->
    -- | Root vertex.
    Int ->
    -- | Tree scanning result from a root vertex.
    VU.Vector a
  scan n tree acc0At toF act root = VU.create $ do
    dp <- VUM.unsafeNew n
    !_ <- foldImpl tree acc0At toF act root $ \v a -> do
      VGM.unsafeWrite dp v a
    pure dp
  
  -- | \(O(n)\) Folds a tree from every vertex, using the rerooting technique.
  --
  -- ==== Constraints
  -- - The action monoid \(f\) must be commutative.
  --
  -- ==== __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
  {-# INLINEABLE foldReroot #-}
  foldReroot ::
    forall w f a.
    (HasCallStack, VU.Unbox w, VU.Unbox a, VU.Unbox f, Monoid f) =>
    -- | The number of vertices.
    Int ->
    -- | Graph as a function.
    (Int -> VU.Vector (Int, w)) ->
    -- | @valAt@:Assignment of initial vertex values.
    (Int -> a) ->
    -- | @toF@: Converts a vertex value into an action onto a neighbor vertex.
    (a -> (Int, w) -> f) ->
    -- | @act@: Performs an action onto a vertex value.
    (f -> a -> a) ->
    -- | Tree folding result from every vertex as a root.
    VU.Vector a
  foldReroot n tree valAt toF act = VU.create $ do
    -- Calculate tree DP for every vertex as a root:
    !dp <- VUM.unsafeNew n
    let reroot parent parentF v1 = do
          let !children = VU.filter ((/= parent) . fst) $ tree v1
          let !fL = VU.scanl' (\ !f (!v2, !w) -> (f <>) . (`toF` (v1, w)) $ treeDp VG.! v2) f0 children
          let !fR = VU.scanr' (\(!v2, !w) !f -> (<> f) . (`toF` (v1, w)) $ treeDp VG.! v2) f0 children
  
          -- save
          let !x1 = (parentF <> VU.last fL) `act` valAt v1
          VGM.unsafeWrite dp v1 x1
  
          VU.iforM_ children $ \i2 (!v2, !w) -> do
            -- composited operator excluding @v2@:
            let !f1 = parentF <> (fL VG.! i2) <> (fR VG.! (i2 + 1))
            let !v1Acc = f1 `act` valAt v1
            let !f2 = toF v1Acc (v2, w)
            reroot v1 f2 v2
  
    reroot (-1 :: Int) f0 root0
    pure dp
    where
      !root0 = 0 :: Int
      !f0 = mempty @f
      !treeDp = scan n tree valAt toF act root0 :: VU.Vector a
  