-- |
-- Module      : Data.Select.Quick
-- Description : Quickselect algorithm on boxed vectors.
-- Copyright   : (c) Donnacha Oisín Kidney, 2018
-- License     : MIT
-- Maintainer  : mail@doisinkidney.com
-- Stability   : experimental
-- Portability : portable
--
-- This module provides an implementation of quickselect on boxed
-- vectors. It has an average time of \(\mathcal{O}(n)\), but a
-- worst-case time of \(\mathcal{O}(n^2)\). For an algorithm with
-- similar performance but a better worst-case time, see
-- "Data.Select.Intro".
module Data.Select.Quick
  (selectBy
  ,select)
  where

import           Data.Vector                (Vector)
import qualified Data.Vector                as Vector
import qualified Data.Vector.Mutable        as MVector

import           Control.Monad.ST

import qualified Data.Select.Mutable.Quick as M

-- | \(\mathcal{O}(n)\). Find the nth item, ordered by the supplied
-- relation.
--
-- prop> i >= 0 && i < length xs ==> sort xs !! i === selectBy (<=) i (Vector.fromList xs)
selectBy :: (a -> a -> Bool) -> Int -> Vector a -> a
selectBy _ i xs
  | i < 0 || i >= Vector.length xs =
      error "Data.Select.Quick.selectBy: index out of bounds."
selectBy lte i xs = runST $ do
    ys <- Vector.thaw xs
    j <- M.select lte ys 0 (Vector.length xs - 1) i
    MVector.unsafeRead ys j
{-# INLINE selectBy #-}

-- | \(\mathcal{O}(n)\). Find the nth smallest item in the vector.
--
-- >>> select 4 (Vector.fromList "this is an example")
-- 'a'
--
-- >>> select 3 (Vector.fromList [0,1,4,2,3,5,6])
-- 3
select :: Ord a => Int -> Vector a -> a
select = selectBy (<=)
{-# INLINE select #-}

-- $setup
-- >>> :set -XFlexibleContexts
-- >>> import Data.List (sort)
-- >>> import Test.QuickCheck