heapsort.hs

-- original: http://d.hatena.ne.jp/kazu-yamamoto/20100528/1275008906
import Control.Applicative
import Control.Monad
import Data.Array.ST
import Data.Array.Unboxed
import Random

----------------------------------------------------------------

type Value = Int

----------------------------------------------------------------

heapSort :: (Ix i, Integral i) =>
            UArray i Value -> UArray i Value
heapSort ua = runSTUArray $ heapsort ua
              
heapsort :: (Ord e, MArray a1 e m, Ix i, Integral i, IArray a e) =>
            a i e -> m (a1 i e)
heapsort ua = do
    let (beg,end) = bounds ua
    sua <- newArray_ (beg,end) -- this sets the type of 's'
    copy ua sua beg end
    forM_ [beg+1..end] $ shiftUp sua
    forM_ [end,end-1..beg+1] $ swapAndShiftDown sua beg
    return sua
  where
    copy from to beg end = forM_ [beg..end] $ \i -> writeArray to i (from ! i)
    swapAndShiftDown arr beg idx = 
        swapAndDo arr beg (\_ _ -> True) idx (shiftDown arr beg (idx - beg))

shiftUp :: (Ord e, MArray a e m, Ix i, Integral i) =>
           a i e -> i -> m ()
shiftUp _ 1 = return ()
shiftUp sua c = swapAndDo sua p (>) c (shiftUp sua p)
  where
    p = c `div` 2

shiftDown :: (Ord e, MArray a e m, Ix i, Num i) =>
             a i e -> i -> i -> m ()
shiftDown sua p n
  | c1 > n    = return ()
  | c1 == n   = swapAndDo sua p (>) c1 (return ())
  | otherwise = do
      let c2 = c1 + 1
      xc1 <- readArray sua c1
      xc2 <- readArray sua c2
      let c = if xc1 > xc2 then c1 else c2
      swapAndDo sua p (>) c (shiftDown sua c n)
  where
    c1 = 2 * p

swapAndDo :: (MArray a e m, Ix i) =>
             a i e -> i -> (e -> e -> Bool) -> i -> m () -> m ()
swapAndDo sua p op c cont = do
    xp <- readArray sua p
    xc <- readArray sua c
    when (xc `op` xp) $ do
        writeArray sua c xp
        writeArray sua p xc
        cont

----------------------------------------------------------------

randomArray :: (IArray a e, Random e, Num e, Ix i, Num i) =>
               i -> e -> IO (a i e)
randomArray n boundary = listArray (1,n) . randomRs (0,boundary) <$> getStdGen

----------------------------------------------------------------

main :: IO ()
main = do
  x <- randomArray 10 100 :: IO (UArray Int Value)
  print x
  print $ heapSort x