Skip to content

Instantly share code, notes, and snippets.

@pnf
Last active May 27, 2017 20:16
Show Gist options
  • Save pnf/6ab57ae1112e72ec5f8a8535674a5ea1 to your computer and use it in GitHub Desktop.
Save pnf/6ab57ae1112e72ec5f8a8535674a5ea1 to your computer and use it in GitHub Desktop.
{-# LANGUAGE RankNTypes #-}
import Data.STRef
import Data.Vector (fromList, toList, freeze, thaw)
import Control.Monad
import Data.Vector.Mutable (MVector, STVector, read, write, swap)
import qualified Data.Vector as V (Vector, length)
import Data.List (sortOn)
import Prelude hiding (read)
import GHC.ST
-- type annotations added to code from http://vaibhavsagar.com/drafts/imperative-haskell.html
type Partition = forall a s . Ord a => MVector s a -> Int -> Int -> ST s Int
partitionFirst :: Partition
partitionFirst array l r = do
p <- read array l
i <- newSTRef (l+1)
forM_ [l+1..(r-1)] $ \j -> do
arrayJ <- read array j
i' <- readSTRef i
when (arrayJ < p) $ do
swap array i' j
modifySTRef' i (+1)
i' <- readSTRef i
swap array (i'-1) l
return (i'-1)
partitionLast :: Partition
partitionLast array l r = do
swap array (r-1) l
partitionFirst array l r
partitionMedian :: Partition
partitionMedian array l r = do
p <- chooseMedian array l r
swap array p l
partitionFirst array l r
chooseMedian :: Partition
chooseMedian array l r = do
h <- read array l
t <- read array (r-1)
let len = r-l
let mid = if (len `mod` 2) == 0
then l + (len `div` 2) - 1
else l + (len `div` 2)
m <- read array mid
let options = sortOn snd [(l, h), (mid, m), (r-1, t)]
return (fst (options !! 1))
quicksort' :: Ord a => STVector s a -> Int -> Int -> Partition -> STRef s Int -> ST s ()
quicksort' array start end partition comparisons = when (start < end) $ do
i <- partition array start end
modifySTRef' comparisons (+ (end-start-1))
quicksort' array start i partition comparisons
quicksort' array (i+1) end partition comparisons
quicksort :: Ord a => V.Vector a -> Partition -> (V.Vector a, Int)
quicksort vector partition = runST $ do
array <- (thaw vector)
comps <- newSTRef 0
quicksort' array 0 (V.length vector) partition comps
res <- freeze array
cnt <- readSTRef comps
return (res,cnt)
qs :: (Ord a) => [a] -> [a]
qs as = toList $ fst $ quicksort (fromList as) partitionFirst
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment