module Data.Vector.Heap
( Heap
, heapify
, sift
, findMin
, deleteMin
, updateMin
, null
, length
) where
import Control.Monad (when)
import Control.Monad.ST
import Data.Bits
import Data.Vector.Array
import Data.Vector.Generic.Mutable as G
import Prelude hiding (length, null)
type Heap s a = MArray s a
heapify :: (G.MVector v a, Ord a) => v s a -> ST s ()
heapify v = go (unsafeShiftR (n2) 1) where
!n = G.length v
go k = when (k >= 0) $ do
sift v k
go (k1)
sift :: (G.MVector v a, Ord a) => v s a -> Int -> ST s ()
sift v root0 = go root0 where
!n = G.length v
go root | child1 <- unsafeShiftL root 1 + 1 = do
when (child1 < n) $ do
r <- G.unsafeRead v root
c1 <- G.unsafeRead v child1
let swap0 = if r > c1 then child1 else root
let child2 = child1 + 1
swap1 <- if child2 < n
then do
s0 <- G.unsafeRead v swap0
c2 <- G.unsafeRead v child2
return $ if s0 > c2 then child2 else swap0
else return swap0
when (swap1 /= root) $ do
s1 <- G.unsafeRead v swap1
G.unsafeWrite v swap1 r
G.unsafeWrite v root s1
go swap1
findMin :: G.MVector v a => v s a -> ST s a
findMin v = G.unsafeRead v 0
deleteMin :: (G.MVector v a, Ord a) => v s a -> ST s (v s a)
deleteMin v = do
let !n = G.length v
a <- G.unsafeRead v (n1)
let !v' = G.unsafeSlice 0 (n1) v
updateMin a v'
return v'
updateMin :: (G.MVector v a, Ord a) => a -> v s a -> ST s ()
updateMin a v = do
G.unsafeWrite v 0 a
sift v 0