quicksort.hs

-- randomized quick sort with 3-way partition
module QuickSort (quicksort) where
import Data.Array.ST
import Control.Monad
import Control.Monad.ST
import System.Random

type QA s = STUArray s Int Int

swap :: Int -> Int -> QA s -> ST s ()
swap i j arr = do
  x <- readArray arr i
  y <- readArray arr j
  writeArray arr i y
  writeArray arr j x

partition :: QA s -> (Int, Int) -> ST s (Int, Int)
partition arr (lo, hi) = do
  x <- readArray arr lo
  (m1, m2) <- foldM (step x) (lo, lo) [succ lo .. hi]
  swap lo m1 arr
  return (m1, m2)
  where
    step = \x (i,j) k -> do
      y <- readArray arr k
      let i' = succ i
          j' = succ j in 
        if y < x then 
          do swap k j' arr; swap i' j' arr; return (i',j')
        else if y == x then
          do swap k j' arr; return (i,j')
        else 
          return (i,j)

quicksortST :: QA s -> (Int, Int) -> StdGen -> ST s ()
quicksortST arr (lo, hi) g = do
  when (lo < hi) $ do
    let (lo', g') = randomR (lo, hi) g
    swap lo lo' arr
    (m1, m2) <- partition arr (lo, hi)
    quicksortST arr (lo, m1 - 1) g'
    quicksortST arr (m2 + 1, hi) g'

quicksort :: [Int] -> [Int]
quicksort [] = []
quicksort l = runST $ do
  arr <- newListArray (1, length l) l :: ST s (QA s)
  bounds <- getBounds arr
  quicksortST arr bounds (mkStdGen 1)
  getElems arr