module Data.Vector.Algorithms.Search
( binarySearch
, binarySearchBy
, binarySearchByBounds
, binarySearchL
, binarySearchLBy
, binarySearchLByBounds
, Comparison
) where
import Prelude hiding (read, length)
import Control.Monad.Primitive
import Data.Bits
import Data.Vector.Generic.Mutable
import Data.Vector.Algorithms.Common (Comparison)
binarySearch :: (PrimMonad m, MVector v e, Ord e)
=> v (PrimState m) e -> e -> m Int
binarySearch = binarySearchBy compare
binarySearchBy :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchBy cmp arr e = binarySearchByBounds cmp arr e 0 (length arr)
binarySearchByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchByBounds cmp arr e l u
| u <= l = return l
| otherwise = do e' <- unsafeRead arr k
case cmp e' e of
LT -> binarySearchByBounds cmp arr e (k+1) u
EQ -> return k
GT -> binarySearchByBounds cmp arr e l k
where k = (u + l) `shiftR` 1
binarySearchL :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int
binarySearchL = binarySearchLBy compare
binarySearchLBy :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchLBy cmp arr e = binarySearchLByBounds cmp arr e 0 (length arr)
binarySearchLByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchLByBounds cmp arr e !l !u
| u <= l = return l
| otherwise = do e' <- unsafeRead arr k
case cmp e' e of
LT -> binarySearchLByBounds cmp arr e (k+1) u
_ -> binarySearchLByBounds cmp arr e l k
where k = (u + l) `shiftR` 1