ac-library-hs-1.0.0.1: src/AtCoder/Internal/String.hs
-- | Internal implementation of @AtCoder.String@ module.
module AtCoder.Internal.String
( -- * Suffix array
saNaive,
saDoubling,
saIsImpl,
saIs,
saIsManual,
)
where
import Control.Monad (unless, when)
import Control.Monad.ST (runST)
import Data.Bit (Bit (..))
import Data.Foldable (for_)
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)
-- TODO: remove `HasCallStack` when we're 100% sure the input is guarded
-- | \(O(n^2)\) Internal implementation of suffix array creation (naive).
--
-- @since 1.0.0.0
{-# INLINE saNaive #-}
saNaive :: (HasCallStack) => VU.Vector Int -> VU.Vector Int
saNaive s =
let n = VU.length s
f l0 r0
| l0 == r0 = GT
| otherwise = inner l0 r0
where
inner l r
| l < n && r < n =
let sl = s VG.! l
sr = s VG.! r
in if sl /= sr
then compare sl sr
else inner (l + 1) (r + 1)
| l == n = LT
| otherwise = GT
in -- modify + generate should fuse
VU.modify (VAI.sortBy f) $ VU.generate n id
-- | \(O(n \log n)\) Internal implementation of suffix array creation (doubling).
--
-- @since 1.0.0.0
{-# INLINE saDoubling #-}
saDoubling :: (HasCallStack) => VU.Vector Int -> VU.Vector Int
saDoubling s = VU.create $ do
let n = VU.length s
sa <- VUM.generate n id
let loop rnk tmp k
| k >= n = pure ()
| otherwise = do
rnk' <- VU.unsafeFreeze rnk
let cmp x y =
let rnkX = rnk' VG.! x
rnkY = rnk' VG.! y
in if rnkX /= rnkY
then compare rnkX rnkY
else
let rx = if x + k < n then rnk' VG.! (x + k) else (-1)
ry = if y + k < n then rnk' VG.! (y + k) else (-1)
in compare rx ry
VAI.sortBy cmp sa
sa' <- VU.unsafeFreeze sa
VGM.write tmp (sa' VG.! 0) 0
VU.zipWithM_
( \saI_ saI -> do
tmpI_ <- VGM.read tmp saI_
VGM.write tmp saI $! tmpI_ + if cmp saI_ saI == LT then 1 else 0
)
sa'
(VU.tail sa')
loop tmp rnk (2 * k)
rnkVec <- VU.thaw s
tmpVec <- VUM.unsafeNew n
loop rnkVec tmpVec 1
pure sa
-- TODO: use backpermute
-- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
--
-- @since 1.0.0.0
{-# INLINE saIsImpl #-}
saIsImpl :: (HasCallStack) => Int -> Int -> VU.Vector Int -> Int -> VU.Vector Int
saIsImpl naiveThreshold doublingThreshold s upper = VU.create $ do
let n = VU.length s
let !ls = VU.create $ do
ls_ <- VUM.unsafeNew @_ @Bit n
VGM.write ls_ (n - 1) $ Bit False
for_ [n - 2, n - 3 .. 0] $ \i -> do
let si = s VG.! i
let si1 = s VG.! (i + 1)
b <-
if si == si1
then VGM.read ls_ (i + 1)
else pure . Bit $ si < si1
VGM.write ls_ i b
pure ls_
let (!sumL, !sumS) = runST $ do
sumL_ <- VUM.replicate (upper + 1) (0 :: Int)
sumS_ <- VUM.replicate (upper + 1) (0 :: Int)
VU.zipWithM_
( \(Bit !b) !si -> do
if not b
then VGM.modify sumS_ (+ 1) si
else VGM.modify sumL_ (+ 1) $ si + 1
)
ls
s
for_ [0 .. upper] $ \i -> do
l <- VGM.read sumL_ i
VGM.modify sumS_ (+ l) i
when (i < upper) $ do
si <- VGM.read sumS_ i
VGM.modify sumL_ (+ si) $ i + 1
sumL' <- VU.unsafeFreeze sumL_
sumS' <- VU.unsafeFreeze sumS_
pure (sumL', sumS')
sa <- VUM.unsafeNew @_ @Int n
let induce lms = do
VGM.set sa (-1)
buf <- VU.thaw sumS
VU.forM_ lms $ \d -> do
unless (d == n) $ do
i <- VGM.read buf $ s VG.! d
VGM.write buf (s VG.! d) $ i + 1
VGM.write sa i d
VU.iforM_ sumL (VGM.write buf)
do
i <- VGM.read buf $ s VG.! (n - 1)
VGM.write buf (s VG.! (n - 1)) $ i + 1
VGM.write sa i $ n - 1
do
-- TODO: try VUM.forM_
for_ [0 .. n - 1] $ \i -> do
v <- VGM.read sa i
when (v >= 1 && not (unBit (ls VG.! (v - 1)))) $ do
j <- VGM.read buf $ s VG.! (v - 1)
VGM.write buf (s VG.! (v - 1)) $ j + 1
VGM.write sa j $ v - 1
VU.iforM_ sumL (VGM.write buf)
-- TODO: try foldr
for_ [n - 1, n - 2 .. 0] $ \i -> do
v <- VGM.read sa i
when (v >= 1 && unBit (ls VG.! (v - 1))) $ do
j <- subtract 1 <$> VGM.read buf (s VG.! (v - 1) + 1)
VGM.write buf (s VG.! (v - 1) + 1) j
VGM.write sa j $ v - 1
let (!lmsMap, !m) = runST $ do
lmsMap_ <- VUM.replicate (n + 1) (-1 :: Int)
len <-
VU.ifoldM'
( \iAcc i (Bit !b1, Bit !b2) -> do
if not b1 && b2
then do
VGM.write lmsMap_ (i + 1) iAcc
pure $ iAcc + 1
else do
pure iAcc
)
(0 :: Int)
$ VU.zip ls (VU.tail ls)
(,len) <$> VU.unsafeFreeze lmsMap_
let lms = VU.create $ do
lms_ <- VUM.unsafeNew @_ @Int m
VU.ifoldM'_
( \iAcc i (Bit !b1, Bit !b2) -> do
if not b1 && b2
then do
VGM.write lms_ iAcc $ i + 1
pure $ iAcc + 1
else do
pure iAcc
)
(0 :: Int)
$ VU.zip ls (VU.tail ls)
pure lms_
induce lms
when (m > 0) $ do
sa' <- VU.unsafeFreeze sa
sortedLms <- VUM.unsafeNew @_ @Int m
VU.foldM'_
( \iAcc v -> do
if lmsMap VG.! v /= -1
then do
VGM.write sortedLms iAcc v
pure $ iAcc + 1
else do
pure iAcc
)
(0 :: Int)
sa'
recS <- VUM.unsafeNew @_ @Int m
(\i -> VGM.write recS i 0) . (lmsMap VG.!) =<< VGM.read sortedLms 0
recUpper' <- do
sortedLms' <- VU.unsafeFreeze sortedLms
VU.foldM'
( \recUpper (!l, !r) -> do
let lmsMapL = lmsMap VG.! l
let lmsMapR = lmsMap VG.! r
let endL = if lmsMapL + 1 < m then lms VG.! (lmsMapL + 1) else n
let endR = if lmsMapR + 1 < m then lms VG.! (lmsMapR + 1) else n
let same
| endL - l /= endR - r = False
| otherwise = inner l r
where
inner x y
| x < endL && s VG.! x == s VG.! y = inner (x + 1) (y + 1)
| x == n || s VG.! x /= s VG.! y = False
| otherwise = True
let recUpper' = if same then recUpper else recUpper + 1
VGM.write recS lmsMapR recUpper'
pure recUpper'
)
(0 :: Int)
$ VU.zip sortedLms' (VU.tail sortedLms')
recS' <- VU.unsafeFreeze recS
let recSa = saIsManual naiveThreshold doublingThreshold recS' recUpper'
VU.iforM_ recSa $ \i x -> do
VGM.write sortedLms i (lms VG.! x)
induce =<< VU.unsafeFreeze sortedLms
pure sa
-- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
--
-- SA-IS, linear-time suffix array construction.
-- Reference:
-- G. Nong, S. Zhang, and W. H. Chan,
-- Two Efficient Algorithms for Linear Time Suffix Array Construction
--
-- @since 1.0.0.0
{-# INLINE saIs #-}
saIs :: (HasCallStack) => VU.Vector Int -> Int -> VU.Vector Int
saIs = saIsManual 10 40
-- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
--
-- SA-IS, linear-time suffix array construction.
-- Reference:
-- G. Nong, S. Zhang, and W. H. Chan,
-- Two Efficient Algorithms for Linear Time Suffix Array Construction
--
-- @since 1.0.0.0
{-# INLINE saIsManual #-}
saIsManual :: (HasCallStack) => Int -> Int -> VU.Vector Int -> Int -> VU.Vector Int
saIsManual naiveThreshold doublingThreshold s upper
| n == 0 = VU.empty
| n == 1 = VU.singleton 0
| n == 2 && s VG.! 0 < s VG.! 1 = VU.fromListN 2 [0, 1]
| n == 2 = VU.fromListN 2 [1, 0]
| n < naiveThreshold = saNaive s
| n < doublingThreshold = saDoubling s
| otherwise = saIsImpl naiveThreshold doublingThreshold s upper
where
n = VU.length s