Binomial Heaps from Okasaki's Purely Functional Data Structures, Chapter 3

pfds3.hs

{-# LANGUAGE IncoherentInstances #-}

import ClassyPrelude hiding (link)
import GHC.TypeLits

--
-- Binomial Trees ala Okasaki's Purely Functional Data Structures, Chap. 3:
--

-- |List whose length is recorded in its type, and also in the type of each element, decreasing from
-- len-1 to 0. Note that a 'SizedList f 1' contains a single element `f 0`, etc.
-- This is a GADT so SNil can be constrained to length 0.
data SizedList (f :: Nat -> *) (len :: Nat) where
  SNil :: SizedList f 0
  SCons :: f l -> SizedList f l -> SizedList f (l + 1)

-- |Binomial tree, which contains a root element and a list of trees of ranks [rank..0].
-- The rank is tracked in the type, so no need to store it as a value.
data BinomialTree a (rank :: Nat) = BinomialTree a (SizedList (BinomialTree a) rank)

instance (Show a) => Show (BinomialTree a rank) where
  show (BinomialTree x ts) = "BT { " <> show x <> "; " <> show ts <> " }"

instance (Show a) => Show (SizedList (BinomialTree a) len) where
  show ts = "[" <> strList ts <> "]"
    where
      strList :: SizedList (BinomialTree a) rank -> String
      strList SNil = ""
      strList (SCons x SNil) = show x
      strList (SCons x xs) = show x <> ", " <> strList xs

one :: a -> BinomialTree a 0
one x = BinomialTree x SNil

-- | The min value is stored in the root node.
root :: BinomialTree a r -> a
root (BinomialTree x _) = x

-- |Link two heap-ordered trees, producing a new heap-ordered tree
link :: (Ord a) => BinomialTree a rank -> BinomialTree a rank -> BinomialTree a (rank + 1)
link t1@(BinomialTree x1 l1) t2@(BinomialTree x2 l2) =
  if x1 <= x2
  then BinomialTree x1 (SCons t2 l1)
  else BinomialTree x2 (SCons t1 l2)

bt1 :: BinomialTree Int 0
bt1 = one 1
bt2 :: BinomialTree Int 1
bt2 = link (one 1) (one 2)
bt3 :: BinomialTree Int 2
bt3 = link (link (one 5) (one 2)) (link (one 4) (one 7))
-- btBad = link (link (one 1) (one 2)) (one 4)


--
-- Binomial Heaps:
--

-- |List of elements of /increasing/ rank; the list's type tracks the rank of the head.
-- Note that the length of the list is unconstrained.
data RankedList (f :: Nat -> *) (rank :: Nat)
  = RNil
  | RCons (f rank) (RankedList f (rank + 1))

-- |Maybe, but tracking rank.
data MaybeF (f :: Nat -> *) (rank :: Nat)
  = FJust (f rank)
  | FNothing
  deriving Show

type BinomialTreeList a rank = RankedList (MaybeF (BinomialTree a)) rank

-- |A Binomial Heap is a (sparse) list of trees of increasing rank.
type BinomialHeap a = BinomialTreeList a 0

instance Show (BinomialTreeList Int rank) where
  show RNil = "[]"
  show (RCons t ts) = show t <> " : " <> show ts

-- |Need a type class to implement insTree, whose behavior depends on the rank of its arguments.
class InsTree r s where
  insTree :: (Ord a) => BinomialTree a s -> BinomialTreeList a r -> BinomialTreeList a r

-- |Insert a tree whose rank is the same as the head of the list's.
instance InsTree r r where
  insTree t RNil = RCons (FJust t) RNil
  insTree t (RCons FNothing ts) = RCons (FJust t) ts
  insTree t (RCons (FJust h) ts) = insTree (link t h) (RCons FNothing ts)

-- |Insert a tree whose rank is greater than the head of the list's.
instance  ((r + 1) ~ s) => InsTree r s where
  insTree t RNil = RCons FNothing (RCons (FJust t) RNil)
  insTree t (RCons FNothing ts) = RCons FNothing (insTree t ts)
  insTree t (RCons (FJust h) ts) = RCons (FJust h) (insTree t ts)

insert :: (Ord a) => a -> BinomialHeap a -> BinomialHeap a
insert x = insTree (one x)

merge :: forall a. (Ord a) => BinomialHeap a -> BinomialHeap a -> BinomialHeap a
merge = loop
  where
    loop :: BinomialTreeList a r -> BinomialTreeList a r -> BinomialTreeList a r
    loop RNil t = t
    loop t RNil = t
    loop (RCons FNothing ts1)   (RCons FNothing ts2)   = RCons FNothing (loop ts1 ts2)
    loop (RCons (FJust h1) ts1) (RCons FNothing ts2)   = RCons (FJust h1) (loop ts1 ts2)
    loop (RCons FNothing ts1)   (RCons (FJust h2) ts2) = RCons (FJust h2) (loop ts1 ts2)
    loop (RCons (FJust h1) ts1) (RCons (FJust h2) ts2) = insTree (link h1 h2) (RCons FNothing (loop ts1 ts2))

-- |Search all the trees, find the one which has the smallest minimum element, and return that
-- tree as well as the heap with that tree removed.
-- Note: both returned in the form of /heaps/, each having the same rank. That means you have to
-- do an 0(log n) search of the "min" list to find the actual value. Could fix that by putting it in
-- a type that just contains one tree, with the rank reified?
removeMinTree :: forall a. (Ord a) => BinomialHeap a -> (BinomialHeap a, BinomialHeap a)
removeMinTree = (\ (_, minH, restH) -> (minH, restH)) . go
  where
    go :: BinomialTreeList a r -> (Maybe a, BinomialTreeList a r, BinomialTreeList a r)
    go RNil                   = (Nothing, RNil, RNil)
    go (RCons FNothing  RNil) = (Nothing, RNil, RNil)
    go (RCons (FJust h) RNil) = (Just (root h), RCons (FJust h) RNil, RNil)
    go (RCons FNothing  ts)   = let (minA, minH, restH) = go ts in (minA, RCons FNothing minH, RCons FNothing restH)
    go (RCons (FJust h) ts)   =
      let (minA, minH, restH) = go ts
      in case minA of
        Just a | a < root h -> (minA, RCons FNothing minH, RCons (FJust h) restH)
        _ ->                   (Just (root h), RCons (FJust h) RNil, RCons FNothing ts)

-- |Find the minimum element by removing its tree and extracting the root.
-- unfortunately complicated by the min tree being wrapped in a heap.
findMin :: forall a. (Ord a) => BinomialHeap a -> Maybe a
findMin = go . fst . removeMinTree
  where
    go :: BinomialTreeList a r -> Maybe a
    go RNil = Nothing
    go (RCons FNothing ts) = go ts
    go (RCons (FJust h) _) = Just (root h)


-- |Ug, parameters in the wrong order for Functor, plus an existential type, oh my.
mapSized :: (forall l. f l -> g l) -> SizedList f len -> SizedList g len
mapSized _ SNil = SNil
mapSized t (SCons x xs) = SCons (t x) (mapSized t xs)

-- Reverse a sized list of any length, to a renked list from 0.
-- Note: This was a lot harder to write than it looks.
reverseSized :: SizedList f len -> RankedList f 0
reverseSized = loop RNil
  where
    loop :: RankedList f r -> SizedList f r -> RankedList f 0
    loop acc SNil = acc
    loop acc (SCons x xs) = loop (RCons x acc) xs


deleteMin :: forall a. (Ord a) => BinomialHeap a -> BinomialHeap a
deleteMin bh = go minT restT
  where
    (minT, restT) = removeMinTree bh

    go :: forall r. BinomialTreeList a r -> BinomialHeap a -> BinomialHeap a
    -- Min tree not present at this rank, so continue:
    go (RCons FNothing minTs) = go minTs
    -- Here's the min tree, so remove the root and merge the rest:
    go (RCons (FJust (BinomialTree _ minTs)) _) = merge (toHeap minTs)
    -- End of the list, reached only if there was no min (empty heap):
    go RNil = id

    -- Convert a naked list of trees to a heap by reversing the list (and wrapping each one in Just).
    -- This seems to be Okasaki's main trick, and it's almost as simple here.
    toHeap = reverseSized . mapSized FJust


-- -- |Ex. 3.5:
-- -- Minimum element by direct recursion over the trees.
-- -- Still obviously O(log n).
findMin' :: (Ord a) => BinomialTreeList a r -> Maybe a
findMin' RNil = Nothing
findMin' (RCons FNothing ts) = findMin' ts
findMin' (RCons (FJust h) ts) = case findMin' ts of
  Nothing -> Just (root h)
  Just minRoot -> Just (min minRoot (root h))


-- -- Ex 3.6
-- data BinomialHeapWithExplicitMin rank a
--   = BinomialHeapWithExplicitMin Int (BinomialHeap a)
--   deriving Show
-- -- The implementation seems trivial


heap5 :: BinomialHeap Int
heap5 = foldr insert RNil [5, 1, 8, 2, 1]


-- something for ghcid to print out:
result = putStrLn . unlines $
  [ tshow $ heap5
  , tshow $ insert 3 heap5
  , tshow $ removeMinTree (insert 3 heap5)
  , tshow $ findMin (insert 3 heap5)
  , tshow $ deleteMin (insert 3 heap5)
  , tshow $ findMin' (insert 3 heap5)
  ]