Quick Sort: Mutable or ST Array

Question

This is the most Monadic code I have written to date :-) and I'd welcome comments.

I'm also struck how much faster this mutable approach is than the immutable version I first wrote (see end). OK, so QuickSort is pretty much on mutability, but most algorithms seems to involve significant manipulation of elements in the data structure and immutability leads to slowness due to the endless copying of e.g. Arrays. I'm left wondering whether immutability is practical in the real world?

-- qsort :: Array -> beginning of subsection -> end of subsection
qsort :: (IOArray Int Int) -> Int -> Int -> IO ()
qsort arr min mx =
    if mx - min < 1 then do
        return ()

    else do
        p <- readArray arr min
        --foldM :: Monad m => (a -> b -> m a) -> a -> [b] -> m a
        final_i <- foldM (partitioner p) (min+1) [(min+1)..mx]
        swap min $ final_i - 1
        qsort arr min     (final_i-2)
        qsort arr final_i mx     

    where
        swap i j = do
            arr_i <- readArray arr i
            arr_j <- readArray arr j
            writeArray arr i arr_j
            writeArray arr j arr_i

        partitioner p i idx = do
            arr_idx <- readArray arr idx
            if arr_idx > p then
                return i
            else do
                swap i idx
                return $ i+1

I have since made the minor modifications to use an ST Monad (Gist). Not sure which is to be considered better Haskell though.

qsort :: (V.Vector Int) -> Int -> Int -> (V.Vector Int)
qsort v i j
    | len < 2 = v
    | j > len - 1 = 
        let
            -- should really switch pivot with last point in < group
            --swap 0 (i-1)
            (less, greater) = V.splitAt i v
            less_sorted = qsort (V.tail less) 1 1
            great_sorted = qsort greater 1 1
        in less_sorted V.++ cons (V.head less) great_sorted
    | vector_j > p = 
        qsort v i (j + 1)
    | otherwise = 
        qsort (swap i j) (i+1) (j+1)
    where
        len = V.length v
        p = v ! 0
        vector_i = v ! i
        vector_j = v ! j
        swap i' j' = update v $ fromList [(i', v!j),(j', v!i)]

Answer 1

Mutability, ST and pure code

Whenever you can use ST instead of IO use ST. ST limits the effects that can happen, and runST (or its variants) make sure that no STRefs (or variants) leave the ST. That is, you can always embed a ST s a where you want to use a. You cannot do the same for IO a.

Unfortunately, update has \$\mathcal O(n+m) \$, and therefore every call of swap is \$\mathcal O(n) \$. That's not a problem when you have a single Vector, e.g.

V.foldr (\x a -> x + 2 * a) 0 . V.map sin . (// foo) . fromList

will generate a single vector due to fusion, not three intermediates. Depending on the used vector version and optimization level, you might even end up with a single loop and no generated vector at all. So immutability doesn't prevent optimizations.

Quick Sort on the other hand really needs mutability. So a ST variant (STUArray or MVector) is more appropriate.

Interface

Now that we know that it's fine to use an ST variant, let's have a look at your code. First we look at qsort's type:

qsort :: (IOArray Int Int) -> Int -> Int -> IO ()

An end user has no idea what Int and Int shall be. They want to sort an array, not ponder the right arguments. Therefore, we add a variant that doesn't use that arguments:

import Control.Monad
import Control.Monad.ST
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M

sort  :: Ord a => V.Vector a -> V.Vector a
sort v = -- left as exercise, should use sort'

sort' :: Ord a => M.STVector s a -> ST s ()
sort' v = qsort v 1 (M.length v)

This interface is easy to use and hard to get wrong. I use the vector variants since I'm more familiar with those, but the same holds for STArray. Note that I relaxed the Int to Ord a => a`.

Internals

Next, we have a look at qsort itself. I'll keep its type for now. First of, instead of

if condition then return () else action

use

when (not condition) action

After all, return does not return prematurely, which can be misleading to Haskell newcomers. Next, try not to shadow Prelude names like min and max. Let's call the arguments lower and upper instead.

Furthermore, you call the value at an index arr_idx. That's misleading. And p is usually used as predicate, so let's call it piv instead. And final_i is just the new midpoint we found with our piv. We end up with.

qsort :: Ord a => (IOArray Int a) -> Int -> Int -> IO ()
qsort arr lower upper =
    when (upper - lower >= 1) $ do
        piv <- readArray arr lower
        midpoint <- foldM (partitioner p) (lower+1) [(lower+1)..upper]
        swap lower $ midpoint - 1
        qsort arr lower   (midpoint - 2)
        qsort arr final_i upper
    where
        swap i j = do
            arr_i <- readArray arr i
            arr_j <- readArray arr j
            writeArray arr i arr_j
            writeArray arr j arr_i

        partitioner piv i idx = do
            value <- readArray arr idx
            if value > piv then
                return i
            else do
                swap i idx
                return $ i+1

Exercise: Try to rewrite qsort with M.Vector or even Data.Vector.Unboxed.Mutable (the latter only needs an additional constraint).

Bottom line

Other than that your mutable code looks fine. I would change swap slightly so that there is only one temporary element, though:

        swap i j = do
            tmp <- readArray arr i
            readArray arr j >>= writeArray arr i
            writeArray arr j tmp

That's personal preference.