module Data.RangeMin.LCA (Index, lowestCommonAncestor, quickLCA) where
import Data.RangeMin.LCA.IndexM
import Control.Monad
import Data.RangeMin
import Data.RangeMin.Common.ST
import Data.RangeMin.Common.Vector
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Primitive as PV
import qualified Data.Vector as V
import Data.Tree
dfOrder :: Tree a -> (Tree (Index, a), Int)
dfOrder tree = execIndexM $ indexer tree where
indexer (Node a ts) = do
i <- getIndex
ts' <- mapM indexer ts
return (Node (i, a) ts')
type Depth = Int
data Trav a = Trav !Depth !Index a
travel :: (a -> Index) -> Tree a -> [Trav a] -> [Trav a]
travel f = trav' 0 where
trav' !d (Node x ts) zs = let me = Trav d (f x) x in
me:foldr (\ t -> trav' (d+1) t . (me:)) zs ts
quickLCA :: Tree a -> (Int, Tree (Index, a), Index -> Index -> (Index, a))
quickLCA tree = case dfOrder tree of
(iTree, n) -> (n, iTree, lowestCommonAncestor n fst iTree)
lowestCommonAncestor :: Int -> (a -> Index) -> Tree a -> Index -> Index -> a
lowestCommonAncestor !n ix tree = vals `seq` lca
where ixs :: UV.Vector Int
depths :: PV.Vector Int
!(!depths, !ixs, !vals) = inlineRunST $ do
let !m = (2 * n 1)
!depthsM <- new m
!ixsM <- new m
!valsM <- new m
let go !i (Trav d j a:ts) = do
write depthsM i d
write ixsM j i
write valsM i a
go (i+1) ts
go _ [] = return ()
go 0 (travel ix tree [])
liftM3 (,,) (unsafeFreeze depthsM) (unsafeFreeze ixsM)
((unsafeFreeze :: V.MVector s a -> ST s (V.Vector a)) valsM)
rM :: Int -> Int -> Int
!rM = intRangeMin depths
lca !i !j = let
iIx = ixs ! i
jIx = ixs ! j
in vals ! case compare iIx jIx of
EQ -> iIx
LT -> rM iIx (jIx iIx + 1)
GT -> rM jIx (iIx jIx + 1)