Skip to content

Instantly share code, notes, and snippets.

@tanakh
Forked from kazu-yamamoto/gist:3051375
Created July 5, 2012 06:43
Show Gist options
  • Save tanakh/3051852 to your computer and use it in GitHub Desktop.
Save tanakh/3051852 to your computer and use it in GitHub Desktop.
QuickSort with STUArray and Vector(s)
{-# LANGUAGE BangPatterns #-}
module Main where
import Control.Applicative
import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.List (sort)
import qualified Data.Vector.Algorithms.Intro as Intro
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Storable as VS
import System.Random
import Criterion.Main
----------------------------------------------------------------
main :: IO ()
main = do
!r1 <- randomList 10000 10000
!r2 <- randomList 100000 100000
defaultMain $ [
bgroup "" [
bench "rnd 10^4 list" $ nf qs r1
, bench "rnd 10^5 list" $ nf qs r2
, bench "rnd 10^4 list (Data.List.sort)" $ nf sort r1
, bench "rnd 10^5 list (Data.List.sort)" $ nf sort r2
, bench "rnd 10^4 STUArray" $ nf quickSort r1
, bench "rnd 10^5 STUArray" $ nf quickSort r2
, bench "rnd 10^4 Vector.Unboxed Intro" $ nf introSortVU r1
, bench "rnd 10^5 Vector.Unboxed Intro" $ nf introSortVU r2
, bench "rnd 10^4 Vector.Storable Intro" $ nf introSortVS r1
, bench "rnd 10^5 Vector.Storable Intro" $ nf introSortVS r2
]]
randomList :: Index -> Value -> IO [Value]
randomList n boundary = replicateM n randomInt
where
randomInt :: IO Value
randomInt = getStdRandom (randomR (0,boundary))
----------------------------------------------------------------
qs :: Ord a => [a] -> [a]
qs [] = []
qs (x:xs) = qs lt ++ [x] ++ qs gt
where
lt = filter (<x) xs
gt = filter (>=x) xs
----------------------------------------------------------------
type Index = Int
type Value = Int
type SUA s = STUArray s Index Value
type PRED = Value -> Value -> Bool
----------------------------------------------------------------
quickSort :: [Value] -> [Value]
quickSort xs = runST $ do
arr <- newListArray (beg,end) xs
quicksort beg end arr
toList beg end arr
where
beg = 0
end = length xs - 1
quicksort :: Index -> Index -> SUA s -> ST s ()
quicksort l u arr
| l >= u = return ()
| otherwise = do
-- swap l random arr -- good pivot
t <- readArray arr l
let i = l
j = u + 1
nj <- loop t u i j arr
swap l nj arr
quicksort l (nj-1) arr
quicksort (nj+1) u arr
loop :: Value -> Index -> Index -> Index -> SUA s -> ST s Index
loop t u i j arr = do
ni <- doWhile i (+1) (< t)
nj <- doWhile j (subtract 1) (> t)
if ni > nj then
return nj
else do
swap ni nj arr
loop t u ni nj arr
where
doWhile k op p
| nk > u = return nk
| otherwise = do
x <- readArray arr nk
if p x then
doWhile nk op p
else
return nk
where
nk = op k
swap :: Index ->Index -> SUA s -> ST s ()
swap i j sua = do
xi <- readArray sua i
xj <- readArray sua j
writeArray sua i xj
writeArray sua j xi
toList :: Index -> Index -> SUA s -> ST s [Value]
toList i lim arr
| i > lim = return []
| otherwise = (:) <$> readArray arr i <*> toList (i+1) lim arr
-----
introSortVU :: [Value] -> [Value]
introSortVU xs = runST $ do
vum <- VU.unsafeThaw $ VU.fromList xs
Intro.sort vum
VU.toList <$> VU.unsafeFreeze vum
introSortVS :: [Value] -> [Value]
introSortVS xs = runST $ do
vum <- VS.unsafeThaw $ VS.fromList xs
Intro.sort vum
VS.toList <$> VS.unsafeFreeze vum
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment