{-# LANGUAGE RecordWildCards #-}

module Data.Graph.Tree.HLD where

import Control.Monad
import Control.Monad.ST
import Data.Function
import qualified Data.Vector.Fusion.Stream.Monadic as MS
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import Data.Graph.Sparse
import My.Prelude ((..<))

type HLDIndex = Int

-- | Heavy-Light-Decomposition
data HLD = HLD
  { HLD -> Vector HLDIndex
indexHLD :: U.Vector HLDIndex
  , HLD -> Vector HLDIndex
parentHLD :: U.Vector Vertex
  , HLD -> Vector HLDIndex
pathHeadHLD :: U.Vector Vertex
  }
  deriving (HLDIndex -> HLD -> ShowS
[HLD] -> ShowS
HLD -> String
(HLDIndex -> HLD -> ShowS)
-> (HLD -> String) -> ([HLD] -> ShowS) -> Show HLD
forall a.
(HLDIndex -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: HLDIndex -> HLD -> ShowS
showsPrec :: HLDIndex -> HLD -> ShowS
$cshow :: HLD -> String
show :: HLD -> String
$cshowList :: [HLD] -> ShowS
showList :: [HLD] -> ShowS
Show)

-- | /O(log V)/
lcaHLD :: HLD -> Vertex -> Vertex -> Vertex
lcaHLD :: HLD -> HLDIndex -> HLDIndex -> HLDIndex
lcaHLD HLD{Vector HLDIndex
indexHLD :: HLD -> Vector HLDIndex
parentHLD :: HLD -> Vector HLDIndex
pathHeadHLD :: HLD -> Vector HLDIndex
indexHLD :: Vector HLDIndex
parentHLD :: Vector HLDIndex
pathHeadHLD :: Vector HLDIndex
..} = HLDIndex -> HLDIndex -> HLDIndex
go
  where
    go :: HLDIndex -> HLDIndex -> HLDIndex
go !HLDIndex
x !HLDIndex
y
      | HLDIndex
ix HLDIndex -> HLDIndex -> Bool
forall a. Ord a => a -> a -> Bool
> HLDIndex
iy = HLDIndex -> HLDIndex -> HLDIndex
go HLDIndex
y HLDIndex
x
      | Bool
otherwise =
          let !hx :: HLDIndex
hx = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
pathHeadHLD HLDIndex
x
              !hy :: HLDIndex
hy = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
pathHeadHLD HLDIndex
y
           in if HLDIndex
hx HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
hy
                then HLDIndex -> HLDIndex -> HLDIndex
go HLDIndex
x (HLDIndex -> HLDIndex) -> HLDIndex -> HLDIndex
forall a b. (a -> b) -> a -> b
$ Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
parentHLD HLDIndex
hy
                else HLDIndex
x
      where
        !ix :: HLDIndex
ix = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
indexHLD HLDIndex
x
        !iy :: HLDIndex
iy = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
indexHLD HLDIndex
y

-- | /O(log V)/
pathHLD :: HLD -> Vertex -> Vertex -> [(HLDIndex, HLDIndex)]
pathHLD :: HLD -> HLDIndex -> HLDIndex -> [(HLDIndex, HLDIndex)]
pathHLD HLD{Vector HLDIndex
indexHLD :: HLD -> Vector HLDIndex
parentHLD :: HLD -> Vector HLDIndex
pathHeadHLD :: HLD -> Vector HLDIndex
indexHLD :: Vector HLDIndex
parentHLD :: Vector HLDIndex
pathHeadHLD :: Vector HLDIndex
..} = HLDIndex -> HLDIndex -> [(HLDIndex, HLDIndex)]
go
  where
    go :: HLDIndex -> HLDIndex -> [(HLDIndex, HLDIndex)]
go !HLDIndex
x !HLDIndex
y
      | HLDIndex
ix HLDIndex -> HLDIndex -> Bool
forall a. Ord a => a -> a -> Bool
> HLDIndex
iy = HLDIndex -> HLDIndex -> [(HLDIndex, HLDIndex)]
go HLDIndex
y HLDIndex
x
      | HLDIndex
hx HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
hy =
          let !ihy :: HLDIndex
ihy = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
indexHLD HLDIndex
hy
              !iy' :: HLDIndex
iy' = HLDIndex
iy HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1
           in (HLDIndex
ihy, HLDIndex
iy') (HLDIndex, HLDIndex)
-> [(HLDIndex, HLDIndex)] -> [(HLDIndex, HLDIndex)]
forall a. a -> [a] -> [a]
: HLDIndex -> HLDIndex -> [(HLDIndex, HLDIndex)]
go HLDIndex
x (Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
parentHLD HLDIndex
hy)
      | HLDIndex
ix HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
== HLDIndex
iy = []
      | Bool
otherwise =
          let !ix' :: HLDIndex
ix' = HLDIndex
ix HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1
              !iy' :: HLDIndex
iy' = HLDIndex
iy HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1
           in [(HLDIndex
ix', HLDIndex
iy')]
      where
        !ix :: HLDIndex
ix = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
indexHLD HLDIndex
x
        !iy :: HLDIndex
iy = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
indexHLD HLDIndex
y
        hx :: HLDIndex
hx = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
pathHeadHLD HLDIndex
x
        hy :: HLDIndex
hy = Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.unsafeIndex Vector HLDIndex
pathHeadHLD HLDIndex
y

-- | /O(V)/
buildHLD :: Vertex -> SparseGraph w -> HLD
buildHLD :: forall w. HLDIndex -> SparseGraph w -> HLD
buildHLD HLDIndex
root gr :: SparseGraph w
gr@SparseGraph{HLDIndex
Vector w
Vector HLDIndex
numVerticesSG :: HLDIndex
numEdgesSG :: HLDIndex
offsetSG :: Vector HLDIndex
adjacentSG :: Vector HLDIndex
edgeCtxSG :: Vector w
numVerticesSG :: forall w. SparseGraph w -> HLDIndex
numEdgesSG :: forall w. SparseGraph w -> HLDIndex
offsetSG :: forall w. SparseGraph w -> Vector HLDIndex
adjacentSG :: forall w. SparseGraph w -> Vector HLDIndex
edgeCtxSG :: forall w. SparseGraph w -> Vector w
..}
  | HLDIndex
numEdgesSG HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
2 HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
* (HLDIndex
numVerticesSG HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
- HLDIndex
1) = String -> HLD
forall a. HasCallStack => String -> a
error String
"not undirected tree"
  | Bool
otherwise = (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
      MVector s HLDIndex
mindexHLD <- HLDIndex -> ST s (MVector (PrimState (ST s)) HLDIndex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HLDIndex -> m (MVector (PrimState m) a)
UM.unsafeNew HLDIndex
numVerticesSG
      MVector s HLDIndex
mparentHLD <- HLDIndex -> HLDIndex -> ST s (MVector (PrimState (ST s)) HLDIndex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HLDIndex -> a -> m (MVector (PrimState m) a)
UM.replicate HLDIndex
numVerticesSG HLDIndex
nothing
      MVector s HLDIndex
mpathHeadHLD <- HLDIndex -> HLDIndex -> ST s (MVector (PrimState (ST s)) HLDIndex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HLDIndex -> a -> m (MVector (PrimState m) a)
UM.replicate HLDIndex
numVerticesSG HLDIndex
nothing

      MVector s HLDIndex
madjacent <- Vector HLDIndex -> ST s (MVector (PrimState (ST s)) HLDIndex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector HLDIndex
adjacentSG
      ST s HLDIndex -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s HLDIndex -> ST s ()) -> ST s HLDIndex -> ST s ()
forall a b. (a -> b) -> a -> b
$
        ((HLDIndex -> HLDIndex -> ST s HLDIndex)
 -> HLDIndex -> HLDIndex -> ST s HLDIndex)
-> HLDIndex -> HLDIndex -> ST s HLDIndex
forall a. (a -> a) -> a
fix
          ( \HLDIndex -> HLDIndex -> ST s HLDIndex
dfs HLDIndex
pv HLDIndex
v -> do
              MVector (PrimState (ST s)) HLDIndex
-> HLDIndex -> HLDIndex -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> HLDIndex -> a -> m ()
UM.write MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
mparentHLD HLDIndex
v HLDIndex
pv
              (HLDIndex
size, (HLDIndex
_, HLDIndex
heavyId)) <-
                ((HLDIndex, (HLDIndex, HLDIndex))
 -> (HLDIndex, HLDIndex) -> ST s (HLDIndex, (HLDIndex, HLDIndex)))
-> (HLDIndex, (HLDIndex, HLDIndex))
-> Vector (HLDIndex, HLDIndex)
-> ST s (HLDIndex, (HLDIndex, HLDIndex))
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
U.foldM'
                  ( \(!HLDIndex
sz, !(HLDIndex, HLDIndex)
mm) (HLDIndex
ei, HLDIndex
nv) -> do
                      HLDIndex
sz' <- HLDIndex -> HLDIndex -> ST s HLDIndex
dfs HLDIndex
v HLDIndex
nv
                      (HLDIndex, (HLDIndex, HLDIndex))
-> ST s (HLDIndex, (HLDIndex, HLDIndex))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (HLDIndex
sz HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
sz', (HLDIndex, HLDIndex)
-> (HLDIndex, HLDIndex) -> (HLDIndex, HLDIndex)
forall a. Ord a => a -> a -> a
max (HLDIndex, HLDIndex)
mm (HLDIndex
sz', HLDIndex
ei))
                  )
                  (HLDIndex
1 :: Int, (HLDIndex
0, HLDIndex
nothing))
                  (Vector (HLDIndex, HLDIndex)
 -> ST s (HLDIndex, (HLDIndex, HLDIndex)))
-> (Vector (HLDIndex, HLDIndex) -> Vector (HLDIndex, HLDIndex))
-> Vector (HLDIndex, HLDIndex)
-> ST s (HLDIndex, (HLDIndex, HLDIndex))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((HLDIndex, HLDIndex) -> Bool)
-> Vector (HLDIndex, HLDIndex) -> Vector (HLDIndex, HLDIndex)
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
U.filter ((HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
pv) (HLDIndex -> Bool)
-> ((HLDIndex, HLDIndex) -> HLDIndex)
-> (HLDIndex, HLDIndex)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HLDIndex, HLDIndex) -> HLDIndex
forall a b. (a, b) -> b
snd)
                  (Vector (HLDIndex, HLDIndex)
 -> ST s (HLDIndex, (HLDIndex, HLDIndex)))
-> Vector (HLDIndex, HLDIndex)
-> ST s (HLDIndex, (HLDIndex, HLDIndex))
forall a b. (a -> b) -> a -> b
$ SparseGraph w
gr SparseGraph w -> HLDIndex -> Vector (HLDIndex, HLDIndex)
forall w. SparseGraph w -> HLDIndex -> Vector (HLDIndex, HLDIndex)
`iadj` HLDIndex
v
              Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (HLDIndex
heavyId HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
nothing) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                MVector (PrimState (ST s)) HLDIndex
-> HLDIndex -> HLDIndex -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> HLDIndex -> HLDIndex -> m ()
UM.swap MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
madjacent HLDIndex
heavyId (Vector HLDIndex
offsetSG Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.! HLDIndex
v)
              HLDIndex -> ST s HLDIndex
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return HLDIndex
size
          )
          HLDIndex
nothing
          HLDIndex
root
      ST s HLDIndex -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s HLDIndex -> ST s ()) -> ST s HLDIndex -> ST s ()
forall a b. (a -> b) -> a -> b
$
        ((HLDIndex -> HLDIndex -> HLDIndex -> HLDIndex -> ST s HLDIndex)
 -> HLDIndex -> HLDIndex -> HLDIndex -> HLDIndex -> ST s HLDIndex)
-> HLDIndex -> HLDIndex -> HLDIndex -> HLDIndex -> ST s HLDIndex
forall a. (a -> a) -> a
fix
          ( \HLDIndex -> HLDIndex -> HLDIndex -> HLDIndex -> ST s HLDIndex
dfs HLDIndex
i HLDIndex
h HLDIndex
pv HLDIndex
v -> do
              MVector (PrimState (ST s)) HLDIndex
-> HLDIndex -> HLDIndex -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> HLDIndex -> a -> m ()
UM.write MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
mindexHLD HLDIndex
v HLDIndex
i
              MVector (PrimState (ST s)) HLDIndex
-> HLDIndex -> HLDIndex -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> HLDIndex -> a -> m ()
UM.write MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
mpathHeadHLD HLDIndex
v HLDIndex
h
              let o :: HLDIndex
o = Vector HLDIndex
offsetSG Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.! HLDIndex
v
              HLDIndex
nv0 <- MVector (PrimState (ST s)) HLDIndex -> HLDIndex -> ST s HLDIndex
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> HLDIndex -> m a
UM.read MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
madjacent HLDIndex
o
              HLDIndex
acc0 <- if HLDIndex
nv0 HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
pv then HLDIndex -> HLDIndex -> HLDIndex -> HLDIndex -> ST s HLDIndex
dfs (HLDIndex
i HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1) HLDIndex
h HLDIndex
v HLDIndex
nv0 else HLDIndex -> ST s HLDIndex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HLDIndex
i
              (HLDIndex -> HLDIndex -> ST s HLDIndex)
-> HLDIndex -> Stream (ST s) HLDIndex -> ST s HLDIndex
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
MS.foldM'
                ( \HLDIndex
acc HLDIndex
j -> do
                    HLDIndex
nv <- MVector (PrimState (ST s)) HLDIndex -> HLDIndex -> ST s HLDIndex
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> HLDIndex -> m a
UM.read MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
madjacent HLDIndex
j
                    if HLDIndex
nv HLDIndex -> HLDIndex -> Bool
forall a. Eq a => a -> a -> Bool
/= HLDIndex
pv
                      then HLDIndex -> HLDIndex -> HLDIndex -> HLDIndex -> ST s HLDIndex
dfs (HLDIndex
acc HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1) HLDIndex
nv HLDIndex
v HLDIndex
nv
                      else HLDIndex -> ST s HLDIndex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HLDIndex
acc
                )
                HLDIndex
acc0
                (Stream (ST s) HLDIndex -> ST s HLDIndex)
-> Stream (ST s) HLDIndex -> ST s HLDIndex
forall a b. (a -> b) -> a -> b
$ (HLDIndex
o HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1) HLDIndex -> HLDIndex -> Stream (ST s) HLDIndex
forall (m :: * -> *).
Monad m =>
HLDIndex -> HLDIndex -> Stream m HLDIndex
..< Vector HLDIndex
offsetSG Vector HLDIndex -> HLDIndex -> HLDIndex
forall a. Unbox a => Vector a -> HLDIndex -> a
U.! (HLDIndex
v HLDIndex -> HLDIndex -> HLDIndex
forall a. Num a => a -> a -> a
+ HLDIndex
1)
          )
          HLDIndex
0
          HLDIndex
root
          HLDIndex
nothing
          HLDIndex
root

      Vector HLDIndex -> Vector HLDIndex -> Vector HLDIndex -> HLD
HLD
        (Vector HLDIndex -> Vector HLDIndex -> Vector HLDIndex -> HLD)
-> ST s (Vector HLDIndex)
-> ST s (Vector HLDIndex -> Vector HLDIndex -> HLD)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) HLDIndex -> ST s (Vector HLDIndex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
mindexHLD
        ST s (Vector HLDIndex -> Vector HLDIndex -> HLD)
-> ST s (Vector HLDIndex) -> ST s (Vector HLDIndex -> 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
<*> MVector (PrimState (ST s)) HLDIndex -> ST s (Vector HLDIndex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
mparentHLD
        ST s (Vector HLDIndex -> HLD) -> ST s (Vector HLDIndex) -> 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
<*> MVector (PrimState (ST s)) HLDIndex -> ST s (Vector HLDIndex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s HLDIndex
MVector (PrimState (ST s)) HLDIndex
mpathHeadHLD
  where
    nothing :: HLDIndex
nothing = -HLDIndex
1