{-# LANGUAGE UndecidableInstances #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Algorithms.BinarySearch
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Algorithms.BinarySearch where

import           Control.Applicative ((<|>))
import           Data.Sequence (Seq, ViewL(..),ViewR(..))
import qualified Data.Sequence as Seq
import           Data.Set (Set)
import qualified Data.Set.Internal as Set
import qualified Data.Vector.Generic as V

--------------------------------------------------------------------------------

-- | Given a monotonic predicate p, a lower bound l, and an upper bound u, with:
--  p l = False
--  p u = True
--  l < u.
--
-- Get the index h such that everything strictly smaller than h has: p i =
-- False, and all i >= h, we have p h = True
--
-- running time: \(O(\log(u - l))\)
{-# SPECIALIZE binarySearch :: (Int -> Bool) -> Int -> Int -> Int #-}
{-# SPECIALIZE binarySearch :: (Word -> Bool) -> Word -> Word -> Word #-}
binarySearch   :: Integral a => (a -> Bool) -> a -> a -> a
binarySearch :: (a -> Bool) -> a -> a -> a
binarySearch a -> Bool
p = a -> a -> a
go
  where
    go :: a -> a -> a
go a
l a
u = let d :: a
d = a
u a -> a -> a
forall a. Num a => a -> a -> a
- a
l
                 m :: a
m = a
l a -> a -> a
forall a. Num a => a -> a -> a
+ (a
d a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2)
             in if a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 then a
u else if a -> Bool
p a
m then a -> a -> a
go a
l a
m
                                             else a -> a -> a
go a
m a
u

-- | Given a value \(\varepsilon\), a monotone predicate \(p\), and two values \(l\) and
-- \(u\) with:
--
-- - \(p l\) = False
-- - \(p u\) = True
-- - \(l < u\)
--
-- we find a value \(h\) such that:
--
-- - \(p h\) = True
-- - \(p (h - \varepsilon)\) = False
--
-- >>> binarySearchUntil (0.1) (>= 0.5) 0 (1 :: Double)
-- 0.5
-- >>> binarySearchUntil (0.1) (>= 0.51) 0 (1 :: Double)
-- 0.5625
-- >>> binarySearchUntil (0.01) (>= 0.51) 0 (1 :: Double)
-- 0.515625
binarySearchUntil       :: (Fractional r, Ord r)
                        => r
                        -> (r -> Bool) -> r -> r -> r
binarySearchUntil :: r -> (r -> Bool) -> r -> r -> r
binarySearchUntil r
eps r -> Bool
p = r -> r -> r
go
  where
    go :: r -> r -> r
go r
l r
u | r
u r -> r -> r
forall a. Num a => a -> a -> a
- r
l r -> r -> Bool
forall a. Ord a => a -> a -> Bool
< r
eps = r
u
           | Bool
otherwise   = let m :: r
m = (r
l r -> r -> r
forall a. Num a => a -> a -> a
+ r
u) r -> r -> r
forall a. Fractional a => a -> a -> a
/ r
2
                           in if r -> Bool
p r
m then r -> r -> r
go r
l r
m else r -> r -> r
go r
m r
u


--------------------------------------------------------------------------------
-- * Binary Searching in some data structure


class BinarySearch v where
  type Index v :: *
  type Elem  v :: *

  -- | Given a monotonic predicate p and a data structure v, find the
  -- element v[h] such that that
  --
  -- for every index i <  h we have p v[i] = False, and
  -- for every inedx i >= h we have p v[i] = True
  --
  -- returns Nothing if no element satisfies p
  --
  -- running time: \(O(T*\log n)\), where \(T\) is the time to execute the
  -- predicate.
  binarySearchIn     :: (Elem v -> Bool) -> v -> Maybe (Elem v)

  -- | Given a monotonic predicate p and a data structure v, find the
  -- index h such that that
  --
  -- for every index i <  h we have p v[i] = False, and
  -- for every inedx i >= h we have p v[i] = True
  --
  -- returns Nothing if no element satisfies p
  --
  -- running time: \(O(T*\log n)\), where \(T\) is the time to execute the
  -- predicate.
  binarySearchIdxIn :: (Elem v -> Bool) -> v -> Maybe (Index v)

--------------------------------------------------------------------------------
-- * Searching on a Sequence

instance BinarySearch (Seq a) where
  type Index (Seq a) = Int
  type Elem  (Seq a) = a

  -- ^ runs in \(O(T*\log^2 n)\) time.
  binarySearchIn :: (Elem (Seq a) -> Bool) -> Seq a -> Maybe (Elem (Seq a))
binarySearchIn Elem (Seq a) -> Bool
p Seq a
s = Seq a -> Int -> a
forall a. Seq a -> Int -> a
Seq.index Seq a
s (Int -> a) -> Maybe Int -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  (Elem (Seq a) -> Bool) -> Seq a -> Maybe (Index (Seq a))
forall v.
BinarySearch v =>
(Elem v -> Bool) -> v -> Maybe (Index v)
binarySearchIdxIn Elem (Seq a) -> Bool
p Seq a
s

  -- ^ runs in \(O(T*\log^2 n)\) time.
  binarySearchIdxIn :: (Elem (Seq a) -> Bool) -> Seq a -> Maybe (Index (Seq a))
binarySearchIdxIn Elem (Seq a) -> Bool
p Seq a
s = case Seq a -> ViewR a
forall a. Seq a -> ViewR a
Seq.viewr Seq a
s of
                            ViewR a
EmptyR                 -> Maybe (Index (Seq a))
forall a. Maybe a
Nothing
                            (Seq a
_ :> a
x)   | Elem (Seq a) -> Bool
p a
Elem (Seq a)
x       -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ case Seq a -> ViewL a
forall a. Seq a -> ViewL a
Seq.viewl Seq a
s of
                              (a
y :< Seq a
_) | Elem (Seq a) -> Bool
p a
Elem (Seq a)
y          -> Int
0
                              ViewL a
_                       -> (Int -> Bool) -> Int -> Int -> Int
forall a. Integral a => (a -> Bool) -> a -> a -> a
binarySearch Int -> Bool
p' Int
0 Int
u
                                       | Bool
otherwise -> Maybe (Index (Seq a))
forall a. Maybe a
Nothing
    where
      p' :: Int -> Bool
p' = a -> Bool
Elem (Seq a) -> Bool
p (a -> Bool) -> (Int -> a) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Seq a -> Int -> a
forall a. Seq a -> Int -> a
Seq.index Seq a
s
      u :: Int
u  = Seq a -> Int
forall a. Seq a -> Int
Seq.length Seq a
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

instance {-# OVERLAPPABLE #-} V.Vector v a => BinarySearch (v a) where
  type Index (v a) = Int
  type Elem  (v a) = a

  binarySearchIdxIn :: (Elem (v a) -> Bool) -> v a -> Maybe (Index (v a))
binarySearchIdxIn Elem (v a) -> Bool
p' v a
v | v a -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
v   = Maybe (Index (v a))
forall a. Maybe a
Nothing
                         | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Bool
p Int
n' = Maybe (Index (v a))
forall a. Maybe a
Nothing
                         | Bool
otherwise  = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ if Int -> Bool
p Int
0 then Int
0 else (Int -> Bool) -> Int -> Int -> Int
forall a. Integral a => (a -> Bool) -> a -> a -> a
binarySearch Int -> Bool
p Int
0 Int
n'
    where
      n' :: Int
n' = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      p :: Int -> Bool
p = a -> Bool
Elem (v a) -> Bool
p' (a -> Bool) -> (Int -> a) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
V.!)

  binarySearchIn :: (Elem (v a) -> Bool) -> v a -> Maybe (Elem (v a))
binarySearchIn Elem (v a) -> Bool
p v a
v = (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
V.!) (Int -> a) -> Maybe Int -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  (Elem (v a) -> Bool) -> v a -> Maybe (Index (v a))
forall v.
BinarySearch v =>
(Elem v -> Bool) -> v -> Maybe (Index v)
binarySearchIdxIn Elem (v a) -> Bool
p v a
v

instance BinarySearch (Set a) where
  type Index (Set a) = Int
  type Elem  (Set a) = a

  binarySearchIn :: (Elem (Set a) -> Bool) -> Set a -> Maybe (Elem (Set a))
binarySearchIn Elem (Set a) -> Bool
p = Set a -> Maybe a
Set a -> Maybe (Elem (Set a))
go
    where
      go :: Set a -> Maybe a
go = \case
        Set a
Set.Tip                     -> Maybe a
forall a. Maybe a
Nothing
        Set.Bin Int
_ a
k Set a
l Set a
r | Elem (Set a) -> Bool
p a
Elem (Set a)
k       -> Set a -> Maybe a
go Set a
l Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> a -> Maybe a
forall a. a -> Maybe a
Just a
k
                        | Bool
otherwise -> Set a -> Maybe a
go Set a
r

  binarySearchIdxIn :: (Elem (Set a) -> Bool) -> Set a -> Maybe (Index (Set a))
binarySearchIdxIn Elem (Set a) -> Bool
p = Set a -> Maybe Int
Set a -> Maybe (Index (Set a))
go
    where
      go :: Set a -> Maybe Int
go = \case
        Set a
Set.Tip                     -> Maybe Int
forall a. Maybe a
Nothing
        Set.Bin Int
_ a
k Set a
l Set a
r | Elem (Set a) -> Bool
p a
Elem (Set a)
k       -> Set a -> Maybe Int
go Set a
l Maybe Int -> Maybe Int -> Maybe Int
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Int -> Maybe Int
forall a. a -> Maybe a
Just (Set a -> Int
forall a. Set a -> Int
Set.size Set a
l)
                        | Bool
otherwise -> (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
Set.size Set a
l)) (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set a -> Maybe Int
go Set a
r