Search.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}

-- ---------------------------------------------------------------------------
-- |
-- Module      : Data.Vector.Algorithms.Search
-- Copyright   : (c) 2009-2010 Dan Doel
-- Maintainer  : Dan Doel <dan.doel@gmail.com>
-- Stability   : Experimental
-- Portability : Non-portable (bang patterns)
--
-- This module implements several methods of searching for indicies to insert
-- elements into a sorted vector.

module Data.Vector.Algorithms.Search
       ( binarySearch
       , binarySearchBy
       , binarySearchByBounds
       , binarySearchL
       , binarySearchLBy
       , binarySearchLByBounds
       , binarySearchR
       , binarySearchRBy
       , binarySearchRByBounds
       , binarySearchP
       , binarySearchPBounds
       , 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)

-- | Finds an index in a given sorted vector at which the given element could
-- be inserted while maintaining the sortedness of the vector.
binarySearch :: (PrimMonad m, MVector v e, Ord e)
             => v (PrimState m) e -> e -> m Int
binarySearch = binarySearchBy compare
{-# INLINE binarySearch #-}

-- | Finds an index in a given vector, which must be sorted with respect to the
-- given comparison function, at which the given element could be inserted while
-- preserving the vector's sortedness.
binarySearchBy :: (PrimMonad m, MVector v e)
               => Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchBy cmp vec e = binarySearchByBounds cmp vec e 0 (length vec)
{-# INLINE binarySearchBy #-}

-- | Given a vector sorted with respect to a given comparison function in indices
-- in [l,u), finds an index in [l,u] at which the given element could be inserted
-- while preserving sortedness.
binarySearchByBounds :: (PrimMonad m, MVector v e)
                     => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchByBounds cmp vec e = loop
 where
 loop !l !u
   | u <= l    = return l
   | otherwise = do e' <- unsafeRead vec k
                    case cmp e' e of
                      LT -> loop (k+1) u
                      EQ -> return k
                      GT -> loop l     k
  where k = (u + l) `shiftR` 1
{-# INLINE binarySearchByBounds #-}

-- | Finds the lowest index in a given sorted vector at which the given element
-- could be inserted while maintaining the sortedness.
binarySearchL :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int
binarySearchL = binarySearchLBy compare
{-# INLINE binarySearchL #-}

-- | Finds the lowest index in a given vector, which must be sorted with respect to
-- the given comparison function, at which the given element could be inserted
-- while preserving the sortedness.
binarySearchLBy :: (PrimMonad m, MVector v e)
                => Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchLBy cmp vec e = binarySearchLByBounds cmp vec e 0 (length vec)
{-# INLINE binarySearchLBy #-}

-- | Given a vector sorted with respect to a given comparison function on indices
-- in [l,u), finds the lowest index in [l,u] at which the given element could be
-- inserted while preserving sortedness.
binarySearchLByBounds :: (PrimMonad m, MVector v e)
                      => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchLByBounds cmp vec e = binarySearchPBounds p vec
 where p e' = case cmp e' e of LT -> False ; _ -> True
{-# INLINE binarySearchLByBounds #-}

-- | Finds the greatest index in a given sorted vector at which the given element
-- could be inserted while maintaining sortedness.
binarySearchR :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int
binarySearchR = binarySearchRBy compare
{-# INLINE binarySearchR #-}

-- | Finds the greatest index in a given vector, which must be sorted with respect to
-- the given comparison function, at which the given element could be inserted
-- while preserving the sortedness.
binarySearchRBy :: (PrimMonad m, MVector v e)
                => Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchRBy cmp vec e = binarySearchRByBounds cmp vec e 0 (length vec)
{-# INLINE binarySearchRBy #-}

-- | Given a vector sorted with respect to the given comparison function on indices
-- in [l,u), finds the greatest index in [l,u] at which the given element could be
-- inserted while preserving sortedness.
binarySearchRByBounds :: (PrimMonad m, MVector v e)
                      => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchRByBounds cmp vec e = binarySearchPBounds p vec
 where p e' = case cmp e' e of GT -> True ; _ -> False
{-# INLINE binarySearchRByBounds #-}

-- | Given a predicate that is guaraneteed to be monotone on the given vector,
-- finds the first index at which the predicate returns True, or the length of
-- the array if the predicate is false for the entire array.
binarySearchP :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int
binarySearchP p vec = binarySearchPBounds p vec 0 (length vec)
{-# INLINE binarySearchP #-}

-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in
-- a given vector, finds the index in [l,u] at which the predicate turns from
-- False to True (yielding u if the entire interval is False).
binarySearchPBounds :: (PrimMonad m, MVector v e)
                    => (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds p vec = loop
 where
 loop !l !u
   | u <= l    = return l
   | otherwise = unsafeRead vec k >>= \e -> if p e then loop l k else loop (k+1) u
  where k = (u + l) `shiftR` 1
{-# INLINE binarySearchPBounds #-}