{-|
== Binary search

Sources:

* https://en.wikipedia.org/wiki/Binary_search_algorithm#Procedure_for_finding_the_leftmost_element

The complexities below assume the predicate takes O(1) time, they should be scaled appropriately if
not so.

-}

{-
Implementation notes:
* binSearch can be written in terms of binSearchM but is more much likely to be used so it is
  separate for ease of copying.
-}

module BinSearch
    ( binSearch
    , binSearchA
    , binSearchM
    ) where

import Data.Array

-- | Returns the first value in [l..h] which satisfies the predicate f, h + 1 otherwise.
-- Returns l if l > h. l + h should not overflow. O(log (h - l + 1)).
binSearch :: Integral i => (i -> Bool) -> i -> i -> i
binSearch :: (i -> Bool) -> i -> i -> i
binSearch i -> Bool
f = i -> i -> i
go where
    go :: i -> i -> i
go i
l i
h
        | i
l i -> i -> Bool
forall a. Ord a => a -> a -> Bool
> i
h     = i
l
        | i -> Bool
f i
m       = i -> i -> i
go i
l (i
m i -> i -> i
forall a. Num a => a -> a -> a
- i
1)
        | Bool
otherwise = i -> i -> i
go (i
m i -> i -> i
forall a. Num a => a -> a -> a
+ i
1) i
h
        where m :: i
m = (i
l i -> i -> i
forall a. Num a => a -> a -> a
+ i
h) i -> i -> i
forall a. Integral a => a -> a -> a
`div` i
2

-- | binSearch on an Array. O(log n).
binSearchA :: (a -> Bool) -> Array Int a -> Int
binSearchA :: (a -> Bool) -> Array Int a -> Int
binSearchA a -> Bool
f Array Int a
a = (Int -> Bool) -> Int -> Int -> Int
forall i. Integral i => (i -> Bool) -> i -> i -> i
binSearch (a -> Bool
f (a -> Bool) -> (Int -> a) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array Int a
aArray Int a -> Int -> a
forall i e. Ix i => Array i e -> i -> e
!)) Int
l Int
h where (Int
l, Int
h) = Array Int a -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds Array Int a
a

-- | binSearch but the predicate returns a monad. O(log (h - l + 1)).
binSearchM :: (Monad m, Integral i) => (i -> m Bool) -> i -> i -> m i
binSearchM :: (i -> m Bool) -> i -> i -> m i
binSearchM i -> m Bool
f = i -> i -> m i
go where
    go :: i -> i -> m i
go i
l i
h | i
l i -> i -> Bool
forall a. Ord a => a -> a -> Bool
> i
h = i -> m i
forall (f :: * -> *) a. Applicative f => a -> f a
pure i
l
    go i
l i
h = do
        let m :: i
m = (i
l i -> i -> i
forall a. Num a => a -> a -> a
+ i
h) i -> i -> i
forall a. Integral a => a -> a -> a
`div` i
2
        Bool
v <- i -> m Bool
f i
m
        if Bool
v then i -> i -> m i
go i
l (i
m i -> i -> i
forall a. Num a => a -> a -> a
- i
1) else i -> i -> m i
go (i
m i -> i -> i
forall a. Num a => a -> a -> a
+ i
1) i
h

--------------------------------------------------------------------------------
-- For tests

-- Allows specialization across modules
{-# INLINABLE binSearch #-}
{-# INLINABLE binSearchA #-}