module Data.PSQueue
(
Binding((:->))
, key
, prio
, PSQ
, size
, null
, lookup
, empty
, singleton
, insert
, insertWith
, delete
, adjust
, adjustWithKey
, update
, updateWithKey
, alter
, keys
, toList
, toAscList
, toDescList
, fromList
, fromAscList
, fromDistinctAscList
, findMin
, deleteMin
, minView
, atMost
, atMostRange
, foldr
, foldl
) where
import Prelude hiding (lookup,null,foldl,foldr)
import qualified Prelude as P
data Binding k p = k :-> p deriving (Eq,Ord,Show,Read)
infix 0 :->
key :: Binding k p -> k
key (k :-> _) = k
prio :: Binding k p -> p
prio (_ :-> p) = p
data PSQ k p = Void | Winner k p (LTree k p) k
instance (Show k, Show p, Ord k, Ord p) => Show (PSQ k p) where
show = show . toAscList
size :: PSQ k p -> Int
size Void = 0
size (Winner _ _ lt _) = 1 + size' lt
null :: PSQ k p -> Bool
null Void = True
null (Winner _ _ _ _) = False
lookup :: (Ord k, Ord p) => k -> PSQ k p -> Maybe p
lookup k q =
case tourView q of
Null -> fail "PSQueue.lookup: Empty queue"
Single k' p
| k == k' -> return p
| otherwise -> fail "PSQueue.lookup: Key not found"
tl `Play` tr
| k <= maxKey tl -> lookup k tl
| otherwise -> lookup k tr
empty :: (Ord k, Ord p) => PSQ k p
empty = Void
singleton :: (Ord k, Ord p) => k -> p -> PSQ k p
singleton k p = Winner k p Start k
insert :: (Ord k, Ord p) => k -> p -> PSQ k p -> PSQ k p
insert k p q =
case tourView q of
Null -> singleton k p
Single k' p' ->
case compare k k' of
LT -> singleton k p `play` singleton k' p'
EQ -> singleton k p
GT -> singleton k' p' `play` singleton k p
tl `Play` tr
| k <= maxKey tl -> insert k p tl `play` tr
| otherwise -> tl `play` insert k p tr
insertWith :: (Ord k, Ord p) => (p->p->p) -> k -> p -> PSQ k p -> PSQ k p
insertWith f = insertWithKey (\_ p p'-> f p p')
insertWithKey :: (Ord k, Ord p) => (k->p->p->p) -> k -> p -> PSQ k p -> PSQ k p
insertWithKey f k p q =
case tourView q of
Null -> singleton k p
Single k' p' ->
case compare k k' of
LT -> singleton k p `play` singleton k' p'
EQ -> singleton k (f k p p')
GT -> singleton k' p' `play` singleton k p
tl `Play` tr
| k <= maxKey tl -> insertWithKey f k p tl `unsafePlay` tr
| otherwise -> tl `unsafePlay` insertWithKey f k p tr
delete :: (Ord k, Ord p) => k -> PSQ k p -> PSQ k p
delete k q =
case tourView q of
Null -> empty
Single k' p
| k == k' -> empty
| otherwise -> singleton k' p
tl `Play` tr
| k <= maxKey tl -> delete k tl `play` tr
| otherwise -> tl `play` delete k tr
adjust :: (Ord p, Ord k) => (p -> p) -> k -> PSQ k p -> PSQ k p
adjust f = adjustWithKey (\_ p -> f p)
adjustWithKey :: (Ord k, Ord p) => (k -> p -> p) -> k -> PSQ k p -> PSQ k p
adjustWithKey f k q =
case tourView q of
Null -> empty
Single k' p
| k == k' -> singleton k' (f k p)
| otherwise -> singleton k' p
tl `Play` tr
| k <= maxKey tl -> adjustWithKey f k tl `unsafePlay` tr
| otherwise -> tl `unsafePlay` adjustWithKey f k tr
update :: (Ord k, Ord p) => (p -> Maybe p) -> k -> PSQ k p -> PSQ k p
update f = updateWithKey (\_ p -> f p)
updateWithKey :: (Ord k, Ord p) => (k -> p -> Maybe p) -> k -> PSQ k p -> PSQ k p
updateWithKey f k q =
case tourView q of
Null -> empty
Single k' p
| k==k' -> case f k p of
Nothing -> empty
Just p' -> singleton k p'
| otherwise -> singleton k' p
tl `Play` tr
| k <= maxKey tl -> updateWithKey f k tl `unsafePlay` tr
| otherwise -> tl `unsafePlay` updateWithKey f k tr
alter :: (Ord k, Ord p) => (Maybe p -> Maybe p) -> k -> PSQ k p -> PSQ k p
alter f k q =
case tourView q of
Null ->
case f Nothing of
Nothing -> empty
Just p -> singleton k p
Single k' p
| k == k' -> case f (Just p) of
Nothing -> empty
Just p' -> singleton k' p'
| otherwise -> case f Nothing of
Nothing -> singleton k' p
Just p' -> insert k p' $ singleton k' p
tl `Play` tr
| k <= maxKey tl -> alter f k tl `unsafePlay` tr
| otherwise -> tl `unsafePlay` alter f k tr
keys :: (Ord k, Ord p) => PSQ k p -> [k]
keys = map key . toList
fromList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromList = P.foldr (\(k:->p) q -> insert k p q) empty
fromAscList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromAscList = fromDistinctAscList . stripEq
where stripEq [] = []
stripEq (x:xs) = stripEq' x xs
stripEq' x' [] = [x']
stripEq' x' (x:xs)
| x' == x = stripEq' x' xs
| otherwise = x' : stripEq' x xs
fromDistinctAscList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromDistinctAscList = foldm unsafePlay empty . map (\(k:->p) -> singleton k p)
foldm :: (a -> a -> a) -> a -> [a] -> a
foldm (*) e x
| P.null x = e
| otherwise = fst (rec (length x) x)
where rec 1 (a : as) = (a, as)
rec n as = (a1 * a2, as2)
where m = n `div` 2
(a1, as1) = rec (n m) as
(a2, as2) = rec m as1
toList :: (Ord k, Ord p) => PSQ k p -> [Binding k p]
toList = toAscList
toAscList :: (Ord k, Ord p) => PSQ k p -> [Binding k p]
toAscList q = seqToList (toAscLists q)
toAscLists :: (Ord k, Ord p) => PSQ k p -> Sequ (Binding k p)
toAscLists q = case tourView q of
Null -> emptySequ
Single k p -> singleSequ (k :-> p)
tl `Play` tr -> toAscLists tl <> toAscLists tr
toDescList :: (Ord k, Ord p) => PSQ k p -> [ Binding k p ]
toDescList q = seqToList (toDescLists q)
toDescLists :: (Ord k, Ord p) => PSQ k p -> Sequ (Binding k p)
toDescLists q = case tourView q of
Null -> emptySequ
Single k p -> singleSequ (k :-> p)
tl `Play` tr -> toDescLists tr <> toDescLists tl
findMin :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p)
findMin Void = Nothing
findMin (Winner k p t m) = Just (k :-> p)
deleteMin :: (Ord k, Ord p) => PSQ k p -> PSQ k p
deleteMin Void = Void
deleteMin (Winner k p t m) = secondBest t m
minView :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p, PSQ k p)
minView Void = Nothing
minView (Winner k p t m) = Just ( k :-> p , secondBest t m )
secondBest :: (Ord k, Ord p) => LTree k p -> k -> PSQ k p
secondBest Start _m = Void
secondBest (LLoser _ k p tl m tr) m' = Winner k p tl m `play` secondBest tr m'
secondBest (RLoser _ k p tl m tr) m' = secondBest tl m `play` Winner k p tr m'
atMost :: (Ord k, Ord p) => p -> PSQ k p -> [Binding k p]
atMost pt q = seqToList (atMosts pt q)
atMosts :: (Ord k, Ord p) => p -> PSQ k p -> Sequ (Binding k p)
atMosts _pt Void = emptySequ
atMosts pt (Winner k p t _) = prune k p t
where
prune k p t
| p > pt = emptySequ
| otherwise = traverse k p t
traverse k p Start = singleSequ (k :-> p)
traverse k p (LLoser _ k' p' tl _m tr) = prune k' p' tl <> traverse k p tr
traverse k p (RLoser _ k' p' tl _m tr) = traverse k p tl <> prune k' p' tr
atMostRange :: (Ord k, Ord p) => p -> (k, k) -> PSQ k p -> [Binding k p]
atMostRange pt (kl, kr) q = seqToList (atMostRanges pt (kl, kr) q)
atMostRanges :: (Ord k, Ord p) => p -> (k, k) -> PSQ k p -> Sequ (Binding k p)
atMostRanges _pt _range Void = emptySequ
atMostRanges pt range@(kl, kr) (Winner k p t _) = prune k p t
where
prune k p t
| p > pt = emptySequ
| otherwise = traverse k p t
traverse k p Start
| k `inrange` range = singleSequ (k :-> p)
| otherwise = emptySequ
traverse k p (LLoser _ k' p' tl m tr) =
guard (kl <= m) (prune k' p' tl) <> guard (m <= kr) (traverse k p tr)
traverse k p (RLoser _ k' p' tl m tr) =
guard (kl <= m) (traverse k p tl) <> guard (m <= kr) (prune k' p' tr)
inrange :: (Ord a) => a -> (a, a) -> Bool
a `inrange` (l, r) = l <= a && a <= r
foldr :: (Ord k,Ord p) => (Binding k p -> b -> b) -> b -> PSQ k p -> b
foldr f z q =
case tourView q of
Null -> z
Single k p -> f (k:->p) z
l`Play`r -> foldr f (foldr f z r) l
foldl :: (Ord k,Ord p) => (b -> Binding k p -> b) -> b -> PSQ k p -> b
foldl f z q =
case tourView q of
Null -> z
Single k p -> f z (k:->p)
l`Play`r -> foldl f (foldl f z l) r
type Size = Int
data LTree k p = Start
| LLoser !Size !k !p (LTree k p) !k (LTree k p)
| RLoser !Size !k !p (LTree k p) !k (LTree k p)
size' :: LTree k p -> Size
size' Start = 0
size' (LLoser s _ _ _ _ _) = s
size' (RLoser s _ _ _ _ _) = s
left, right :: LTree a b -> LTree a b
left Start = error "left: empty loser tree"
left (LLoser _ _ _ tl _ _ ) = tl
left (RLoser _ _ _ tl _ _ ) = tl
right Start = error "right: empty loser tree"
right (LLoser _ _ _ _ _ tr) = tr
right (RLoser _ _ _ _ _ tr) = tr
maxKey :: PSQ k p -> k
maxKey Void = error "maxKey: empty queue"
maxKey (Winner _k _p _t m) = m
lloser, rloser :: k -> p -> LTree k p -> k -> LTree k p -> LTree k p
lloser k p tl m tr = LLoser (1 + size' tl + size' tr) k p tl m tr
rloser k p tl m tr = RLoser (1 + size' tl + size' tr) k p tl m tr
omega :: Int
omega = 4
lbalance, rbalance ::
(Ord k, Ord p) => k-> p -> LTree k p -> k -> LTree k p -> LTree k p
lbalance k p l m r
| size' l + size' r < 2 = lloser k p l m r
| size' r > omega * size' l = lbalanceLeft k p l m r
| size' l > omega * size' r = lbalanceRight k p l m r
| otherwise = lloser k p l m r
rbalance k p l m r
| size' l + size' r < 2 = rloser k p l m r
| size' r > omega * size' l = rbalanceLeft k p l m r
| size' l > omega * size' r = rbalanceRight k p l m r
| otherwise = rloser k p l m r
lbalanceLeft k p l m r
| size' (left r) < size' (right r) = lsingleLeft k p l m r
| otherwise = ldoubleLeft k p l m r
lbalanceRight k p l m r
| size' (left l) > size' (right l) = lsingleRight k p l m r
| otherwise = ldoubleRight k p l m r
rbalanceLeft k p l m r
| size' (left r) < size' (right r) = rsingleLeft k p l m r
| otherwise = rdoubleLeft k p l m r
rbalanceRight k p l m r
| size' (left l) > size' (right l) = rsingleRight k p l m r
| otherwise = rdoubleRight k p l m r
lsingleLeft k1 p1 t1 m1 (LLoser _ k2 p2 t2 m2 t3)
| p1 <= p2 = lloser k1 p1 (rloser k2 p2 t1 m1 t2) m2 t3
| otherwise = lloser k2 p2 (lloser k1 p1 t1 m1 t2) m2 t3
lsingleLeft k1 p1 t1 m1 (RLoser _ k2 p2 t2 m2 t3) =
rloser k2 p2 (lloser k1 p1 t1 m1 t2) m2 t3
rsingleLeft k1 p1 t1 m1 (LLoser _ k2 p2 t2 m2 t3) =
rloser k1 p1 (rloser k2 p2 t1 m1 t2) m2 t3
rsingleLeft k1 p1 t1 m1 (RLoser _ k2 p2 t2 m2 t3) =
rloser k2 p2 (rloser k1 p1 t1 m1 t2) m2 t3
lsingleRight k1 p1 (LLoser _ k2 p2 t1 m1 t2) m2 t3 =
lloser k2 p2 t1 m1 (lloser k1 p1 t2 m2 t3)
lsingleRight k1 p1 (RLoser _ k2 p2 t1 m1 t2) m2 t3 =
lloser k1 p1 t1 m1 (lloser k2 p2 t2 m2 t3)
rsingleRight k1 p1 (LLoser _ k2 p2 t1 m1 t2) m2 t3 =
lloser k2 p2 t1 m1 (rloser k1 p1 t2 m2 t3)
rsingleRight k1 p1 (RLoser _ k2 p2 t1 m1 t2) m2 t3
| p1 <= p2 = rloser k1 p1 t1 m1 (lloser k2 p2 t2 m2 t3)
| otherwise = rloser k2 p2 t1 m1 (rloser k1 p1 t2 m2 t3)
ldoubleLeft k1 p1 t1 m1 (LLoser _ k2 p2 t2 m2 t3) =
lsingleLeft k1 p1 t1 m1 (lsingleRight k2 p2 t2 m2 t3)
ldoubleLeft k1 p1 t1 m1 (RLoser _ k2 p2 t2 m2 t3) =
lsingleLeft k1 p1 t1 m1 (rsingleRight k2 p2 t2 m2 t3)
ldoubleRight k1 p1 (LLoser _ k2 p2 t1 m1 t2) m2 t3 =
lsingleRight k1 p1 (lsingleLeft k2 p2 t1 m1 t2) m2 t3
ldoubleRight k1 p1 (RLoser _ k2 p2 t1 m1 t2) m2 t3 =
lsingleRight k1 p1 (rsingleLeft k2 p2 t1 m1 t2) m2 t3
rdoubleLeft k1 p1 t1 m1 (LLoser _ k2 p2 t2 m2 t3) =
rsingleLeft k1 p1 t1 m1 (lsingleRight k2 p2 t2 m2 t3)
rdoubleLeft k1 p1 t1 m1 (RLoser _ k2 p2 t2 m2 t3) =
rsingleLeft k1 p1 t1 m1 (rsingleRight k2 p2 t2 m2 t3)
rdoubleRight k1 p1 (LLoser _ k2 p2 t1 m1 t2) m2 t3 =
rsingleRight k1 p1 (lsingleLeft k2 p2 t1 m1 t2) m2 t3
rdoubleRight k1 p1 (RLoser _ k2 p2 t1 m1 t2) m2 t3 =
rsingleRight k1 p1 (rsingleLeft k2 p2 t1 m1 t2) m2 t3
play :: (Ord k, Ord p) => PSQ k p -> PSQ k p -> PSQ k p
Void `play` t' = t'
t `play` Void = t
Winner k p t m `play` Winner k' p' t' m'
| p <= p' = Winner k p (rbalance k' p' t m t') m'
| otherwise = Winner k' p' (lbalance k p t m t') m'
unsafePlay :: (Ord k, Ord p) => PSQ k p -> PSQ k p -> PSQ k p
Void `unsafePlay` t' = t'
t `unsafePlay` Void = t
Winner k p t m `unsafePlay` Winner k' p' t' m'
| p <= p' = Winner k p (rbalance k' p' t m t') m'
| otherwise = Winner k' p' (lbalance k p t m t') m'
data TourView k p = Null | Single k p | PSQ k p `Play` PSQ k p
tourView :: (Ord k) => PSQ k p -> TourView k p
tourView Void = Null
tourView (Winner k p Start _m) = Single k p
tourView (Winner k p (RLoser _ k' p' tl m tr) m') =
Winner k p tl m `Play` Winner k' p' tr m'
tourView (Winner k p (LLoser _ k' p' tl m tr) m') =
Winner k' p' tl m `Play` Winner k p tr m'
emptySequ :: Sequ a
singleSequ :: a -> Sequ a
(<>) :: Sequ a -> Sequ a -> Sequ a
seqFromList :: [a] -> Sequ a
seqFromListT :: ([a] -> [a]) -> Sequ a
seqToList :: Sequ a -> [a]
infixr 5 <>
newtype Sequ a = Sequ ([a] -> [a])
emptySequ = Sequ (\as -> as)
singleSequ a = Sequ (\as -> a : as)
Sequ x1 <> Sequ x2 = Sequ (\as -> x1 (x2 as))
seqFromList as = Sequ (\as' -> as ++ as')
seqFromListT as = Sequ as
seqToList (Sequ x) = x []
instance Show a => Show (Sequ a) where
showsPrec d a = showsPrec d (seqToList a)
guard :: Bool -> Sequ a -> Sequ a
guard False _as = emptySequ
guard True as = as