okasaki

Okasaki's Purely Functional Data Structures
git clone git://git.jtobin.io/okasaki.git
Log | Files | Refs | LICENSE

Binomial.hs (2810B)

      1 {-# OPTIONS_GHC -Wall -fno-warn-unused-top-binds #-}
      2 {-# LANGUAGE TemplateHaskell #-}
      3 
      4 module Okasaki.Heap.Binomial where
      5 
      6 import Data.Eq.Deriving (deriveEq1)
      7 import Data.Fix hiding (cata, ana, hylo)
      8 import Data.Foldable (toList)
      9 import Data.Functor.Foldable
     10 import Data.List (sortOn)
     11 import Data.List.NonEmpty (NonEmpty(..))
     12 import Data.Maybe (fromMaybe)
     13 import Text.Show.Deriving
     14 
     15 data TreeF a r = NodeF a [r]
     16   deriving (Eq, Show, Functor, Foldable, Traversable)
     17 
     18 type Tree a = Fix (TreeF a)
     19 
     20 $(deriveShow1 ''TreeF)
     21 $(deriveEq1 ''TreeF)
     22 
     23 -- exercise 3.6: remove redundant rank info
     24 data Bush a = Bush {-# UNPACK #-} !Int !(Tree a)
     25   deriving Show
     26 
     27 newtype Heap a = Heap [Bush a]
     28   deriving Show
     29 
     30 lef :: Heap a
     31 lef = Heap []
     32 
     33 nod :: a -> [Tree a] -> Tree a
     34 nod a t = Fix (NodeF a t)
     35 
     36 lin :: Ord a => Bush a -> Bush a -> Bush a
     37 lin (Bush i s@(project -> NodeF a c)) (Bush j t@(project -> NodeF b d))
     38   | i /= j    = error "(okasaki, binomial): internal error"
     39   | a <= b    = Bush (succ i) (nod a (t : c))
     40   | otherwise = Bush (succ i) (nod b (s : d))
     41 
     42 ran :: Bush a -> Int
     43 ran (Bush i _) = i
     44 
     45 roo :: Bush a -> a
     46 roo (Bush _ (project -> NodeF a _)) = a
     47 
     48 pet :: Ord a => Bush a -> [Bush a] -> [Bush a]
     49 pet a as = toList $ cata alg as (a :| as)
     50   where
     51     alg l n@(t :| ts) = case l of
     52       Nil -> n
     53       Cons h f
     54         | ran t < ran h -> n
     55         | otherwise     -> f $! lin t h :| drop 1 ts
     56 
     57 put :: Ord a => a -> Heap a -> Heap a
     58 put a (Heap as) = Heap $
     59   pet (Bush 0 (nod a [])) as
     60 
     61 mer :: Ord a => Heap a -> Heap a -> Heap a
     62 mer (Heap a) (Heap b) = Heap $ mor a b where
     63   mor l r = case (l, r) of
     64     (p, []) -> p
     65     ([], q) -> q
     66     (h : t, i : s)
     67       | ran h < ran i -> h : mor t r
     68       | ran i < ran h -> i : mor l s
     69       | otherwise     -> pet (lin h i) (mor t s)
     70 
     71 out :: Ord a => Heap a -> Maybe (Bush a, Heap a)
     72 out (Heap as) = case as of
     73   []       -> Nothing
     74   (h : []) -> Just (h, lef)
     75   (h : t)  -> do
     76     (i, Heap s) <- out (Heap t)
     77     pure $ if   roo h <= roo i
     78            then (h, Heap t)
     79            else (i, Heap (h:s))
     80 
     81 bot :: Ord a => Heap a -> Maybe a
     82 bot h = do
     83   (a, _) <- out h
     84   pure $ roo a
     85 
     86 -- exercise 3.5: direct bot
     87 bed :: Ord a => Heap a -> Maybe a
     88 bed (Heap a) = case sortOn roo a of
     89   []      -> Nothing
     90   (h : _) -> Just (roo h)
     91 
     92 cut :: Ord a => Heap a -> Heap a
     93 cut h = fromMaybe lef $ do
     94   (Bush i (project -> NodeF _ s), t) <- out h
     95   let p = fmap (Bush (pred i)) s
     96   pure $ mer (Heap (reverse p)) t
     97 
     98 -- reference
     99 
    100 -- arbitrary recursion
    101 pyt :: Ord a => Bush a -> Heap a -> Heap a
    102 pyt a (Heap as) = Heap (loop a as) where
    103   loop s ts = case ts of
    104     [] -> pure s
    105     (h : t)
    106       | ran s < ran h -> s : ts
    107       | otherwise     -> loop (lin s h) t
    108 
    109 -- test
    110 
    111 test0 :: Heap Int
    112 test0 = put 3 . put 2 . put 1 . put 0 $ lef
    113 
    114 test1 :: Heap Int
    115 test1 = put 4 test0