Heap sort using State monad along with lenses to deal with heap.

heapsort.hs

{-# LANGUAGE TemplateHaskell, Rank2Types #-}
module Main where

import Prelude hiding (last)

import Control.Lens
import Control.Monad.State.Strict (StateT, evalStateT, put)
import Control.Monad.IO.Class (liftIO)
import Control.Monad (when)
import Control.Applicative ((<$>), (<*>), pure)

import Data.List (intercalate)
import Data.Maybe (fromJust)

import System.IO (hFlush, stdout)

-- | Heap is either empty heap or a node.
type Heap a = Maybe (HeapNode a)

-- | Heap node.
data HeapNode a = HeapNode
    { _val   :: !a          -- ^ Value stored in the node.
    , _left  :: !(Heap a)   -- ^ Left subheap.
    , _right :: !(Heap a)   -- ^ Right subheap.
    }
    deriving (Show)
makeLenses ''HeapNode

type H v = StateT (Heap v) IO
type N v = StateT (HeapNode v) IO

-- | Pointer to the root value.
val_ :: Traversal' (Heap a) a
val_ = traverse . val

-- | Pointers to subheaps.
left_ :: Traversal' (HeapNode a) (HeapNode a)
left_ = left . traverse

-- | Pointer to the right subheap.
right_ :: Traversal' (HeapNode a) (HeapNode a)
right_ = right . traverse

-- | Lens to a leaf.
leaf :: Lens' (Heap a) (Heap a)
leaf f Nothing                               = f Nothing
leaf f h@(Just (HeapNode _ Nothing Nothing)) = f h
leaf f h@(Just (HeapNode x (Just l) r)) = Just <$> (flip (HeapNode x) r <$> (leaf f $ Just l))
leaf f h@(Just (HeapNode x l (Just r))) = Just <$> (HeapNode x l <$> (leaf f $ Just r))

-- | Delete all nodes from the heap, displaying sorted values.
printSorted :: (Show a, Ord a) => H a ()
printSorted = do
    v <- delete
    case v of
        Nothing -> return ()
        Just x  -> liftIO $ putStr (show x ++ " ")

-- | Remove root from a Heap.
delete :: (Ord a) => H a (Maybe a)
delete = do
    x <- preuse val_    -- extract root value
    y <- zoom leaf $ do
        z <- preuse val_    -- extract leaf value
        put Nothing         -- remove leaf
        return z
    val_ .= fromJust y  -- replace root value with leaf's
    down                -- recover consistency
    return x

-- | Swap root and given node values.
swap :: Traversal' (HeapNode a) (HeapNode a) -> N a ()
swap node_ = do
    x <- use val    -- extract root value
    y <- zoom (singular node_) $ do
        z <- use val    -- extract node value
        val .= x        -- replace node value with root's
        return z
    val .= y    -- replace root value with node's

-- | Recover consistency in a Heap.
down :: (Ord a) => H a ()
down = zoom traverse $ do
    zoom left  down     -- go down for the left subheap
    zoom right down     -- go down for the right subheap
    down'

down' :: (Ord a) => N a ()
down' = do
    -- extract root and children values
    [l, r, v] <- mapM preuse [left_.val, right_.val, val]
    let m = maximum [l, v, r]
    
    -- if root value is not maximum
    when (m > v) $ do
        let node_ :: Traversal' (HeapNode a) (HeapNode a)
            node_ = if m > l then right_ else left_
        swap node_          -- swap values with max child
        zoom node_ down'    -- repeat for that child

-- | Convert list to an arbitrary Heap.
listToHeap :: [a] -> Heap a
listToHeap [] = Nothing
listToHeap (x:xs) = Just $ HeapNode x l r
    where
        (l, r) = both %~ listToHeap $ splitAt (1 + length xs `div` 2) xs

-- | Sort list using intermediate Heap structure.
heapSort :: (Show a, Ord a) => [a] -> IO ()
heapSort = evalStateT (down >> printSorted) . listToHeap

-- | Ask user for an input.
prompt :: String -> IO String
prompt p = do
    putStr p
    hFlush stdout
    getLine

main :: IO ()
main = do
    -- read N
    n   <- read <$> prompt "n  = "
    -- read N values
    xs  <- (map read . take n . words) <$> prompt "xs = "
    -- heap sort
    heapSort (xs :: [Int])

This works ~100 times slower than C program (https://gist.github.com/iley/5543784) and ~10 times slower than Java program (https://gist.github.com/pertu/f4550148e21ddee92481).