-- | It contains string algorithms.
--
-- Let \(s\) be a string. We denote the substring of \(s\) between \(a\)-th and \(b - 1\)-th
-- character by \(s[a..b)\).
--
-- ==== __Examples__
--
-- ===== Suffix Array and LCP Array
--
-- >>> import AtCoder.String qualified as S
-- >>> import Data.ByteString.Char8 qualified as BS
-- >>> let s = BS.pack "aab"
-- >>> let sa = S.suffixArrayBS s
-- >>> S.lcpArrayBS s sa
-- [1,0]
--
-- ===== Z Algorithm
--
-- >>> import AtCoder.String qualified as S
-- >>> import Data.ByteString.Char8 qualified as BS
-- >>> let s = BS.pack "abab"
-- >>> S.zAlgorithmBS s
-- [4,0,2,0]
--
-- @since 1.0.0.0
module AtCoder.String
  ( -- * Suffix array
    suffixArray,
    suffixArrayBS,
    suffixArrayOrd,

    -- * LCP array
    lcpArray,
    lcpArrayBS,

    -- * Z algorithm
    zAlgorithm,
    zAlgorithmBS,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.String qualified as ACIS
import Control.Monad.ST (runST)
import Data.ByteString.Char8 qualified as BS
import Data.Char (ord)
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: document `chr`

-- | Calculates suffix array for a `Int` vector.
--
-- Given a string \(s\) of length \(n\), it returns the suffix array of \(s\). Here, the suffix array
-- \(\mathrm{sa}\) of \(s\) is a permutation of \(0, \cdots, n-1\) such that \(s[\mathrm{sa}[i]..n) < s[\mathrm{sa}[i+1]..n)\)
-- holds for all \(i = 0,1, \cdots ,n-2\).
--
-- ==== Constraints
-- - \(0 \leq n\)
-- - \(0 \leq \mathrm{upper} \leq 10^8\)
-- - \(0 \leq x \leq \mathrm{upper}\) for all elements \(x\) of \(s\).
--
-- ==== Complexity
-- - \(O(n + \mathrm{upper})\)-time
--
-- @since 1.0.0.0
{-# INLINE suffixArray #-}
suffixArray :: (HasCallStack) => VU.Vector Int -> Int -> VU.Vector Int
suffixArray :: HasCallStack => Vector Int -> Int -> Vector Int
suffixArray Vector Int
s Int
upper =
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
upper) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.String.suffixArray: given negative `upper`: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
upper
      !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert ((Int -> Bool) -> Vector Int -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all (\Int
d -> Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
d Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
upper) Vector Int
s) String
"AtCoder.String.suffixArray: some input out of bounds"
   in HasCallStack => Vector Int -> Int -> Vector Int
Vector Int -> Int -> Vector Int
ACIS.saIs Vector Int
s Int
upper

-- | Calculates suffix array for a @ByteString@.
--
-- ==== Constraints
-- - \(0 \leq n\)
--
-- ==== Complexity
-- - \(O(n)\)-time
--
-- @since 1.0.0.0
{-# INLINE suffixArrayBS #-}
suffixArrayBS :: (HasCallStack) => BS.ByteString -> VU.Vector Int
suffixArrayBS :: HasCallStack => ByteString -> Vector Int
suffixArrayBS ByteString
s = do
  let n :: Int
n = ByteString -> Int
BS.length ByteString
s
      s2 :: Vector Int
s2 = (Char -> Int) -> Vector Char -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map Char -> Int
ord (Vector Char -> Vector Int) -> Vector Char -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> String -> Vector Char
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
n (ByteString -> String
BS.unpack ByteString
s)
   in HasCallStack => Vector Int -> Int -> Vector Int
Vector Int -> Int -> Vector Int
ACIS.saIs Vector Int
s2 Int
255

-- | Calculates suffix array for a `Ord` type vector.
--
-- ==== Constraints
-- - \(0 \leq n\)
--
-- ==== Complexity
-- - \(O(n \log n)\)-time
-- - \(O(n)\)-space
--
-- @since 1.0.0.0
{-# INLINE suffixArrayOrd #-}
suffixArrayOrd :: (HasCallStack, Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int
suffixArrayOrd :: forall a. (HasCallStack, Ord a, Unbox a) => Vector a -> Vector Int
suffixArrayOrd Vector a
s =
  let n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
s
      (!Int
upper, !Vector Int
s2) = (forall s. ST s (Int, Vector Int)) -> (Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Int, Vector Int)) -> (Int, Vector Int))
-> (forall s. ST s (Int, Vector Int)) -> (Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
        let f :: Int -> Int -> Ordering
f Int
i Int
j = a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i) (Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
j)
        -- modify + generate should fuse
        let idx :: Vector Int
idx = (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify ((Int -> Int -> Ordering)
-> MVector (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
VAI.sortBy Int -> Int -> Ordering
f) (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n Int -> Int
forall a. a -> a
id
        MVector s Int
vec <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
        Int
upper_ <-
          (Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s Int
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
            ( \Int
now Int
i -> do
                let now' :: Int
now' =
                      if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Vector Int
idx Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Vector Int
idx Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i)
                        then Int
now Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                        else Int
now
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
vec (Vector Int
idx Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i) Int
now'
                Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
now'
            )
            (Int
0 :: Int)
            (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n Int -> Int
forall a. a -> a
id)
        (Int
upper_,) (Vector Int -> (Int, Vector Int))
-> ST s (Vector Int) -> ST s (Int, Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
vec
   in HasCallStack => Vector Int -> Int -> Vector Int
Vector Int -> Int -> Vector Int
ACIS.saIs Vector Int
s2 Int
upper

-- | Given a string \(s\) of length \(n\), it returns the LCP array of \(s\). Here, the LCP array of
-- \(s\) is the array of length \(n-1\), such that the \(i\)-th element is the length of the LCP
-- (Longest Common Prefix) of \(s[\mathrm{sa}[i]..n)\) and \(s[\mathrm{sa}[i+1]..n)\).
--
-- ==== Constraints
-- - The second argument is the suffix array of \(s\).
-- - \(1 \leq n\)
--
-- ==== Complexity
-- - \(O(n)\)
--
-- @since 1.0.0.0
{-# INLINE lcpArray #-}
lcpArray ::
  (HasCallStack, Ord a, VU.Unbox a) =>
  -- | A vector representing a string
  VU.Vector a ->
  -- | Suffix array
  VU.Vector Int ->
  -- | LCP array
  VU.Vector Int
lcpArray :: forall a.
(HasCallStack, Ord a, Unbox a) =>
Vector a -> Vector Int -> Vector Int
lcpArray Vector a
s Vector Int
sa =
  let n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
s
      !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1) String
"AtCoder.String.lcpArray: given empty input"
      rnk :: Vector Int
rnk = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
rnkVec <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Int Int
n
        Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
sa ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
saI -> do
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
rnkVec Int
saI Int
i
        MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
rnkVec
   in (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
lcp <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        (Int -> Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> Int -> b -> m a) -> a -> Vector b -> m ()
VU.ifoldM'_
          ( \ !Int
h_ Int
i Int
rnkI -> do
              let h :: Int
h = if Int
h_ Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Int
h_ Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else Int
h_
              if Int
rnkI Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
h
                else do
                  let j :: Int
j = Vector Int
sa Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
rnkI Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                  let inner :: Int -> Int
inner !Int
h'
                        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Bool -> Bool -> Bool
&& Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Int
h'
                        | Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h') a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h') = Int
h'
                        | Bool
otherwise = Int -> Int
inner (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
h' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                  let !h' :: Int
h' = Int -> Int
inner Int
h
                  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
lcp (Int
rnkI Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
h'
                  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
h'
          )
          (Int
0 :: Int)
          Vector Int
rnk
        MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
lcp

-- | @ByteString@ variant of `lcpArray`.
--
-- ==== Constraints
-- - The second argument is the suffix array of \(s\).
-- - \(1 \leq n\)
--
-- ==== Complexity
-- - \(O(n)\)
--
-- @since 1.0.0.0
{-# INLINE lcpArrayBS #-}
lcpArrayBS ::
  (HasCallStack) =>
  -- | String
  BS.ByteString ->
  -- | Suffix array
  VU.Vector Int ->
  -- | LCP array
  VU.Vector Int
lcpArrayBS :: HasCallStack => ByteString -> Vector Int -> Vector Int
lcpArrayBS ByteString
s Vector Int
sa =
  let n :: Int
n = ByteString -> Int
BS.length ByteString
s
      s2 :: Vector Int
s2 = (Char -> Int) -> Vector Char -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map Char -> Int
ord (Vector Char -> Vector Int)
-> (String -> Vector Char) -> String -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String -> Vector Char
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
n (String -> Vector Int) -> String -> Vector Int
forall a b. (a -> b) -> a -> b
$ ByteString -> String
BS.unpack ByteString
s
   in Vector Int -> Vector Int -> Vector Int
forall a.
(HasCallStack, Ord a, Unbox a) =>
Vector a -> Vector Int -> Vector Int
lcpArray Vector Int
s2 Vector Int
sa

-- | Given a `Ord` vector of length \(n\), it returns the array of length \(n\), such that the
-- \(i\)-th element is the length of the LCP (Longest Common Prefix) of \(s[0..n)\) and \(s[i..n)\).
--
-- ==== Constraints
-- - \(n \leq n\)
--
-- ==== Complexity
-- - \(O(n)\)
--
-- @since 1.0.0.0
{-# INLINE zAlgorithm #-}
zAlgorithm :: (Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int
zAlgorithm :: forall a. (Ord a, Unbox a) => Vector a -> Vector Int
zAlgorithm Vector a
s
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector Int
forall a. Unbox a => Vector a
VU.empty
  | Bool
otherwise = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
      MVector s Int
z <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Int Int
n
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
z Int
0 Int
0
      (Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m ()
VU.foldM'_
        ( \Int
j Int
i -> do
            Int
zj <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
z Int
j
            Int
k0 <-
              if Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
zj Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i
                then Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
                else do
                  Int
zij <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
z (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j)
                  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
zj Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int
zij
            let loop :: Int -> Int
loop Int
k
                  | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Bool -> Bool -> Bool
&& Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a
s Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) = Int -> Int
loop (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                  | Bool
otherwise = Int
k
            let k :: Int
k = Int -> Int
loop Int
k0
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
z Int
i Int
k
            Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$
              if Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
zj Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k
                then Int
i
                else Int
j
        )
        (Int
0 :: Int)
        (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
z Int
0 Int
n
      MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
z
  where
    n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
s

-- | Given a string of length \(n\), it returns the array of length \(n\), such that the \(i\)-th
-- element is the length of the LCP (Longest Common Prefix) of \(s[0..n)\) and \(s[i..n)\).
--
-- ==== Constraints
-- - \(n \leq n\)
--
-- ==== Complexity
-- - \(O(n)\)
--
-- @since 1.0.0.0
{-# INLINE zAlgorithmBS #-}
zAlgorithmBS :: BS.ByteString -> VU.Vector Int
zAlgorithmBS :: ByteString -> Vector Int
zAlgorithmBS ByteString
s = Vector Char -> Vector Int
forall a. (Ord a, Unbox a) => Vector a -> Vector Int
zAlgorithm (Vector Char -> Vector Int) -> Vector Char -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> String -> Vector Char
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN (ByteString -> Int
BS.length ByteString
s) (ByteString -> String
BS.unpack ByteString
s)