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