PriorityQueue.hs

module PriorityQueue (
    PQ,
    size, priority, peek,
    singleton, branch, push,
    pop
  ) where

{- Priority queue with ordered keys `k` and values `v`. A branch stores the
   minimum key in its subtrees as well as the value associated (for fast peek)
   as well as its size (for easily maintaining balance). -}

data PQ k v = Leaf k v
            | Branch k v Int (PQ k v) (PQ k v)

{- Reading -}

-- | Extract the size of a given priority queue: O(1)
size :: PQ v a -> Int
size (Leaf _ _) = 1
size (Branch _ _ n _ _) = n

-- | Extract the minimum priority in a given priority queue: O(1)
priority :: PQ v a -> v
priority (Leaf p _) = p
priority (Branch p _ _ _ _) = p

-- | Get the minimum priority element in a given priority queue: O(1)
peek :: PQ v a -> a
peek (Leaf _ e) = e
peek (Branch _ e _ _ _) = e


{- Introducing -}

-- | Construct priority queue with a single element (of the given priority and
-- value): O(1)
singleton :: v -> a -> PQ v a
singleton = Leaf

-- | Safe constructor: creates a branch by merging the taggings of the subtrees
-- according to the semigroup instance.
branch :: Ord k => PQ k v -> PQ k v -> PQ k v
branch x y
 | priority x <  priority y = Branch (priority x) (peek x) (size x + size y) x y
 | priority x >= priority y = Branch (priority y) (peek y) (size x + size y) x y

-- | Given a priority and a value, insert these into a priority queue. This 
-- maintains the balanced property by always inserting into the smaller of
-- the subtrees: O(log n)
push :: Ord v => v -> a -> PQ v a -> PQ v a
push pri val pq@(Leaf _ _) = branch (singleton pri val) pq
push pri val pq@(Branch _ _ _ l r)
  | size l > size r = let r' = push pri val r in branch l r'
  | otherwise       = let l' = push pri val l in branch l' r

{- Eliminating -}

-- | Given a priority queue of size > 1, return the priority queue obtained
-- after removing the element with the lowest priority: O(log n)
pop :: Ord v => PQ v a -> PQ v a
pop (Leaf _ _) = error "Priority queue must have size at least 1"
pop pq@(Branch _ _ _ l r)
  | priority pq == priority l = if isLeaf l then r else branch (pop l) r
  | priority pq == priority r = if isLeaf r then l else branch l (pop r)
  where
    isLeaf :: PQ v a -> Bool
    isLeaf (Leaf _ _) = True
    isLeaf _          = False