1{-# LANGUAGE BangPatterns #-} 2{-# LANGUAGE TypeFamilies #-} 3 4-- --------------------------------------------------------------------------- 5-- | 6-- Module : Data.Vector.Algorithms.Search 7-- Copyright : (c) 2009-2015 Dan Doel, 2015 Tim Baumann 8-- Maintainer : Dan Doel <dan.doel@gmail.com> 9-- Stability : Experimental 10-- Portability : Non-portable (bang patterns) 11-- 12-- This module implements several methods of searching for indicies to insert 13-- elements into a sorted vector. 14 15module Data.Vector.Algorithms.Search 16 ( binarySearch 17 , binarySearchBy 18 , binarySearchByBounds 19 , binarySearchL 20 , binarySearchLBy 21 , binarySearchLByBounds 22 , binarySearchR 23 , binarySearchRBy 24 , binarySearchRByBounds 25 , binarySearchP 26 , binarySearchPBounds 27 , gallopingSearchLeftP 28 , gallopingSearchLeftPBounds 29 , gallopingSearchRightP 30 , gallopingSearchRightPBounds 31 , Comparison 32 ) where 33 34import Prelude hiding (read, length) 35 36import Control.Monad.Primitive 37 38import Data.Bits 39 40import Data.Vector.Generic.Mutable 41 42import Data.Vector.Algorithms.Common (Comparison, midPoint) 43 44-- | Finds an index in a given sorted vector at which the given element could 45-- be inserted while maintaining the sortedness of the vector. 46binarySearch :: (PrimMonad m, MVector v e, Ord e) 47 => v (PrimState m) e -> e -> m Int 48binarySearch = binarySearchBy compare 49{-# INLINE binarySearch #-} 50 51-- | Finds an index in a given vector, which must be sorted with respect to the 52-- given comparison function, at which the given element could be inserted while 53-- preserving the vector's sortedness. 54binarySearchBy :: (PrimMonad m, MVector v e) 55 => Comparison e -> v (PrimState m) e -> e -> m Int 56binarySearchBy cmp vec e = binarySearchByBounds cmp vec e 0 (length vec) 57{-# INLINE binarySearchBy #-} 58 59-- | Given a vector sorted with respect to a given comparison function in indices 60-- in [l,u), finds an index in [l,u] at which the given element could be inserted 61-- while preserving sortedness. 62binarySearchByBounds :: (PrimMonad m, MVector v e) 63 => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int 64binarySearchByBounds cmp vec e = loop 65 where 66 loop !l !u 67 | u <= l = return l 68 | otherwise = do e' <- unsafeRead vec k 69 case cmp e' e of 70 LT -> loop (k+1) u 71 EQ -> return k 72 GT -> loop l k 73 where k = midPoint u l 74{-# INLINE binarySearchByBounds #-} 75 76-- | Finds the lowest index in a given sorted vector at which the given element 77-- could be inserted while maintaining the sortedness. 78binarySearchL :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int 79binarySearchL = binarySearchLBy compare 80{-# INLINE binarySearchL #-} 81 82-- | Finds the lowest index in a given vector, which must be sorted with respect to 83-- the given comparison function, at which the given element could be inserted 84-- while preserving the sortedness. 85binarySearchLBy :: (PrimMonad m, MVector v e) 86 => Comparison e -> v (PrimState m) e -> e -> m Int 87binarySearchLBy cmp vec e = binarySearchLByBounds cmp vec e 0 (length vec) 88{-# INLINE binarySearchLBy #-} 89 90-- | Given a vector sorted with respect to a given comparison function on indices 91-- in [l,u), finds the lowest index in [l,u] at which the given element could be 92-- inserted while preserving sortedness. 93binarySearchLByBounds :: (PrimMonad m, MVector v e) 94 => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int 95binarySearchLByBounds cmp vec e = binarySearchPBounds p vec 96 where p e' = case cmp e' e of LT -> False ; _ -> True 97{-# INLINE binarySearchLByBounds #-} 98 99-- | Finds the greatest index in a given sorted vector at which the given element 100-- could be inserted while maintaining sortedness. 101binarySearchR :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int 102binarySearchR = binarySearchRBy compare 103{-# INLINE binarySearchR #-} 104 105-- | Finds the greatest index in a given vector, which must be sorted with respect to 106-- the given comparison function, at which the given element could be inserted 107-- while preserving the sortedness. 108binarySearchRBy :: (PrimMonad m, MVector v e) 109 => Comparison e -> v (PrimState m) e -> e -> m Int 110binarySearchRBy cmp vec e = binarySearchRByBounds cmp vec e 0 (length vec) 111{-# INLINE binarySearchRBy #-} 112 113-- | Given a vector sorted with respect to the given comparison function on indices 114-- in [l,u), finds the greatest index in [l,u] at which the given element could be 115-- inserted while preserving sortedness. 116binarySearchRByBounds :: (PrimMonad m, MVector v e) 117 => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int 118binarySearchRByBounds cmp vec e = binarySearchPBounds p vec 119 where p e' = case cmp e' e of GT -> True ; _ -> False 120{-# INLINE binarySearchRByBounds #-} 121 122-- | Given a predicate that is guaraneteed to be monotone on the given vector, 123-- finds the first index at which the predicate returns True, or the length of 124-- the array if the predicate is false for the entire array. 125binarySearchP :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int 126binarySearchP p vec = binarySearchPBounds p vec 0 (length vec) 127{-# INLINE binarySearchP #-} 128 129-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in 130-- a given vector, finds the index in [l,u] at which the predicate turns from 131-- False to True (yielding u if the entire interval is False). 132binarySearchPBounds :: (PrimMonad m, MVector v e) 133 => (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int 134binarySearchPBounds p vec = loop 135 where 136 loop !l !u 137 | u <= l = return l 138 | otherwise = unsafeRead vec k >>= \e -> if p e then loop l k else loop (k+1) u 139 where k = midPoint u l 140{-# INLINE binarySearchPBounds #-} 141 142-- | Given a predicate that is guaranteed to be monotone on the vector elements 143-- in order, finds the index at which the predicate turns from False to True. 144-- The length of the vector is returned if the predicate is False for the entire 145-- vector. 146-- 147-- Begins searching at the start of the vector, in increasing steps of size 2^n. 148gallopingSearchLeftP 149 :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int 150gallopingSearchLeftP p vec = gallopingSearchLeftPBounds p vec 0 (length vec) 151{-# INLINE gallopingSearchLeftP #-} 152 153-- | Given a predicate that is guaranteed to be monotone on the vector elements 154-- in order, finds the index at which the predicate turns from False to True. 155-- The length of the vector is returned if the predicate is False for the entire 156-- vector. 157-- 158-- Begins searching at the end of the vector, in increasing steps of size 2^n. 159gallopingSearchRightP 160 :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int 161gallopingSearchRightP p vec = gallopingSearchRightPBounds p vec 0 (length vec) 162{-# INLINE gallopingSearchRightP #-} 163 164-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in 165-- a given vector, finds the index in [l,u] at which the predicate turns from 166-- False to True (yielding u if the entire interval is False). 167-- Begins searching at l, going right in increasing (2^n)-steps. 168gallopingSearchLeftPBounds :: (PrimMonad m, MVector v e) 169 => (e -> Bool) 170 -> v (PrimState m) e 171 -> Int -- ^ l 172 -> Int -- ^ u 173 -> m Int 174gallopingSearchLeftPBounds p vec l u 175 | u <= l = return l 176 | otherwise = do x <- unsafeRead vec l 177 if p x then return l else iter (l+1) l 2 178 where 179 binSearch = binarySearchPBounds p vec 180 iter !i !j !_stepSize | i >= u - 1 = do 181 x <- unsafeRead vec (u-1) 182 if p x then binSearch (j+1) (u-1) else return u 183 iter !i !j !stepSize = do 184 x <- unsafeRead vec i 185 if p x then binSearch (j+1) i else iter (i+stepSize) i (2*stepSize) 186{-# INLINE gallopingSearchLeftPBounds #-} 187 188-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in 189-- a given vector, finds the index in [l,u] at which the predicate turns from 190-- False to True (yielding u if the entire interval is False). 191-- Begins searching at u, going left in increasing (2^n)-steps. 192gallopingSearchRightPBounds :: (PrimMonad m, MVector v e) 193 => (e -> Bool) 194 -> v (PrimState m) e 195 -> Int -- ^ l 196 -> Int -- ^ u 197 -> m Int 198gallopingSearchRightPBounds p vec l u 199 | u <= l = return l 200 | otherwise = iter (u-1) (u-1) (-1) 201 where 202 binSearch = binarySearchPBounds p vec 203 iter !i !j !_stepSize | i <= l = do 204 x <- unsafeRead vec l 205 if p x then return l else binSearch (l+1) j 206 iter !i !j !stepSize = do 207 x <- unsafeRead vec i 208 if p x then iter (i+stepSize) i (2*stepSize) else binSearch (i+1) j 209{-# INLINE gallopingSearchRightPBounds #-} 210