Ticket #3909: QuickBinom.hs

File QuickBinom.hs, 4.1 KB (added by LouisWasserman, 4 years ago)

A quickie implementation of a binomial heap than handles sparse binomial forests well. Not noticeably better than earlier typesafe version, even with all the obvious optimizations I could think of, so I prefer the above, typesafe, elegant version.

Line 
1{-# LANGUAGE CPP #-}
2
3module Data.PQueue.QuickBinom (insert, extract, fromList, toAscList)  where
4
5import Data.List(foldl')
6#ifdef __GLASGOW_HASKELL__
7import GHC.Exts (build)
8import Text.Read (Lexeme(Ident), lexP, parens, prec,
9        readPrec, readListPrec, readListPrecDefault)
10#endif
11
12data PQueue a = Empty | PQueue {-# UNPACK #-} !Int (BinomQueue a)
13-- type PQueue = BinomQueue
14data BinomQueue a = Nil | Cons {-# UNPACK #-} !Int {-# UNPACK #-} !(BinHeap a) (BinomQueue a)
15data BinomQueue' a = Nil' | Cons' {-# UNPACK #-} !(BinHeap a) (BinomQueue' a)
16data BinHeap a = Bin a (BinomQueue' a)
17
18data Extract a = NoExtract | YesExtract a {-# UNPACK #-} !Int (BinomQueue' a) (BinomQueue a)
19
20revQueue :: Int -> BinomQueue' a -> BinomQueue a
21revQueue k = revQueue' (k-1) Nil where
22        revQueue' :: Int -> BinomQueue a -> BinomQueue' a -> BinomQueue a
23        revQueue' k rev ts = k `seq` case ts of
24                Cons' t ts      -> revQueue' (k-1) (Cons k t rev) ts
25                Nil'            -> rev
26
27tip :: a -> BinHeap a
28tip x = Bin x Nil'
29
30data ViewQ a = EmptyQ           -- ^ empty queue
31        | a :^ PQueue a         -- ^ the top (minimum) of the queue and the rest of the queue
32--      deriving (Eq, Ord, Read, Show)
33
34extract :: Ord a => PQueue a -> Maybe (a, PQueue a)
35extract Empty = Nothing
36extract (PQueue n h) = case extractQueue (<=) h of
37        NoExtract       -> Nothing
38        YesExtract x' rk ts tss
39                -> Just (x', PQueue (n-1) (merge (<=) (revQueue rk ts) tss))
40
41insert :: Ord a => a -> PQueue a -> PQueue a
42insert x Empty = PQueue 1 (Cons 0 (tip x) Nil)
43insert x (PQueue n h) = PQueue (n+1) (ins x h)
44        where ins = carry1 (<=) 0 . tip
45
46fromList :: Ord a => [a] -> PQueue a
47fromList = foldr insert Empty
48
49{-# INLINE toAscList #-}
50-- | /O(n log n)/.  Extracts the elements of the priority queue in ascending order.
51toAscList :: Ord a => PQueue a -> [a]
52#ifdef __GLASGOW_HASKELL__
53toAscList q = build (\ c nil -> foldrAsc c nil q)
54#else
55toAscList = foldrAsc (:) []
56#endif
57
58foldrAsc :: Ord a => (a -> b -> b) -> b -> PQueue a -> b
59foldrAsc f z (PQueue n h) = foldrHeap (<=) f z h
60foldrAsc _ z _ = z
61
62foldrHeap :: (a -> a -> Bool) -> (a -> b -> b) -> b -> BinomQueue a -> b
63foldrHeap (<=) f z = foldrH' where
64        foldrH' h = case extractQueue (<=) h of
65                NoExtract       -> z
66                YesExtract x rk ts tss
67                        -> x `f` foldrH' (merge (<=) (revQueue rk ts) tss)
68
69{-# NOINLINE extractQueue #-}
70extractQueue :: (a -> a -> Bool) -> BinomQueue a -> Extract a
71extractQueue _ Nil = NoExtract
72extractQueue (<=) (Cons k t@(Bin x ts) tss) = case extractQueue (<=) tss of
73        NoExtract       -> YesExtract x k ts tss
74        YesExtract minK rk minKids rest
75                | x <= minK     -> YesExtract x k ts tss
76                | otherwise     -> YesExtract minK rk minKids (Cons k t rest)
77
78meldH :: (a -> a -> Bool) -> Int -> BinHeap a -> BinHeap a -> BinHeap a
79meldH (<=) k t1@(Bin x1 ts1) t2@(Bin x2 ts2)
80        | k `seq` x1 <= x2      = Bin x1 (Cons' t2 ts1)
81        | otherwise             = Bin x2 (Cons' t1 ts2)
82
83merge :: (a -> a -> Bool) -> BinomQueue a -> BinomQueue a -> BinomQueue a
84merge (<=) ts1 ts2 = case (ts1, ts2) of
85        (Nil, _)        -> ts2
86        (_, Nil)        -> ts1
87        (Cons k1 t1 ts1', Cons k2 t2 ts2') -> case compare k1 k2 of
88                LT      -> Cons k1 t1 (merge (<=) ts1' ts2)
89                EQ      -> carry (<=) (k1 + 1) (meld k1 t1 t2) ts1' ts2'
90                GT      -> Cons k2 t2 (merge (<=) ts1 ts2')
91        where   meld = meldH (<=)
92
93-- Invariant: k0 <= rank of first trees in ts1, ts2
94carry :: (a -> a -> Bool) -> Int -> BinHeap a -> BinomQueue a -> BinomQueue a -> BinomQueue a
95carry (<=) k0 t0 ts1 ts2 = k0 `seq` t0 `seq` case (ts1, ts2) of
96        (Nil, _)        -> carry1 (<=) k0 t0 ts2
97        (_, Nil)        -> carry1 (<=) k0 t0 ts1
98        (Cons k1 t1 ts1', Cons k2 t2 ts2') -> case (k0 == k1, k0 == k2) of
99                (True, True)    -> Cons k0 t0 (carry (<=) (k0+1) (meld k0 t1 t2) ts1' ts2')
100                (True, False)   -> carry (<=) (k0+1) (meld k0 t0 t1) ts1' ts2
101                (False, True)   -> carry (<=) (k0+1) (meld k0 t0 t2) ts1 ts2'
102                (False, False)  -> Cons k0 t0 $ case compare k1 k2 of
103                        LT      -> Cons k1 t1 (merge (<=) ts1' ts2)
104                        EQ      -> carry (<=) (k1 + 1) (meld k1 t1 t2) ts1' ts2'
105                        GT      -> Cons k2 t2 (merge (<=) ts1 ts2')
106        where   meld = meldH (<=)
107
108carry1 :: (a -> a -> Bool) -> Int -> BinHeap a -> BinomQueue a -> BinomQueue a
109carry1 (<=) k0 t0 ts = k0 `seq` t0 `seq` case ts of
110        Cons k t ts'
111                | k0 == k       -> carry1 (<=) (k0 + 1) (meld k0 t0 t) ts'
112        _                       -> Cons k0 t0 ts
113        where   meld = meldH (<=)