select.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

import           Control.Monad.Primitive
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector.Algorithms.Insertion as Insertion

import Control.Monad
import qualified Test.QuickCheck.Arbitrary as QC
import qualified Test.QuickCheck.Gen as QC

select :: (UV.Unbox a, Ord a, PrimMonad m)
       => Int
       -> UM.MVector (PrimState m) a
       -> m a
select k xs
  | UM.length xs <= 5 = do Insertion.sort xs
                           UM.unsafeRead xs k
  | otherwise = do

      let xs' = xs

      -- calculate the right amount of subsequences
      let
        n :: Int
        n = divUp (UM.length xs') 5

      -- 1. slice the input into sequences of length 5
      -- 2. sort the sequences with insertion sort (inplace)
      -- 3. Take the median
      ys  <- UV.generateM n (\i -> do -- slicing is O(1)
                                      let j = i * 5
                                          m = min 5 (UM.length xs' - j)
                                          subSeq = UM.unsafeSlice j m xs'
                                      -- length(part) <= 5 -> sort is O(1)
                                      Insertion.sort subSeq
                                      -- determine the median O(1)
                                      median subSeq
                           )
      ys' <- UV.unsafeThaw ys

      -- calculate the right amount of subsequences for the medians
      let
        n' :: Int
        n' = divUp (UM.length xs') 10

      a <- select n' ys'

      -- Quicksort partition
      (less, equal, greater) <- partition a xs'

      let
        nls :: Int
        nls = UM.length less

        neq :: Int
        neq = UM.length equal

      if k > nls && k <= (nls + neq)
        then return a
        else if k <= nls
             then select k less
             else select (k - nls - neq) greater
        where
          -- an uprounding div
          -- Such that 5 `divUp` 5 == 1,
          -- and 6 `divUp` 5 == 2
          divUp :: Int -> Int -> Int
          divUp x y = (x + y - 1) `div` y

median :: (UV.Unbox a, PrimMonad m)
       => UM.MVector (PrimState m) a
       -> m a
median v = UM.unsafeRead v (UM.length v `div` 2)

partition :: (UV.Unbox a, Ord a, PrimMonad m)
          => a
          -> UM.MVector (PrimState m) a
          -> m ( UM.MVector (PrimState m) a
               , UM.MVector (PrimState m) a
               , UM.MVector (PrimState m) a
               )
partition a xs = do
  j <- GM.unstablePartition (< a) xs
  let (l, pr) = UM.splitAt j xs
  k <- GM.unstablePartition ( == a) pr
  let (m, r) = UM.splitAt k xs
  return (l, m, r)

main :: IO ()
main = do
  !v <- UV.fromList <$> QC.generate (QC.resize 100000 ((replicateM 100000 (QC.suchThat QC.arbitrary (> 0)))))
  putStrLn "gen"

  v' <- UV.unsafeThaw v
  a <- select 1 v'
  print (a :: Int)