{-# LANGUAGE ScopedTypeVariables #-}
module Data.DEPQ (
DEPQ,
empty,
fromList, toList,
null,
valid,
size,
insert, delete, deleteMin, deleteMax, popMin, popMax,
lookup, findMin, findMax,
topK, bottomK
) where
import Data.Maybe (fromMaybe)
import Data.Ord (Down(..))
import qualified Data.Sequence as S (Seq, empty, (|>))
import Control.DeepSeq (NFData (rnf))
import qualified Data.IntPSQ as P (IntPSQ, empty, null, size, insert, delete, toList, findMin, delete, deleteMin, valid, lookup)
import Prelude hiding (null, lookup)
import Test.QuickCheck (Arbitrary(..), Gen)
data DEPQ p a = DEPQ {
DEPQ p a -> IntPSQ p a
minHeap :: P.IntPSQ p a
, DEPQ p a -> IntPSQ (Down p) a
maxHeap :: P.IntPSQ (Down p) a
} deriving (DEPQ p a -> DEPQ p a -> Bool
(DEPQ p a -> DEPQ p a -> Bool)
-> (DEPQ p a -> DEPQ p a -> Bool) -> Eq (DEPQ p a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall p a. (Ord p, Eq a) => DEPQ p a -> DEPQ p a -> Bool
/= :: DEPQ p a -> DEPQ p a -> Bool
$c/= :: forall p a. (Ord p, Eq a) => DEPQ p a -> DEPQ p a -> Bool
== :: DEPQ p a -> DEPQ p a -> Bool
$c== :: forall p a. (Ord p, Eq a) => DEPQ p a -> DEPQ p a -> Bool
Eq, Int -> DEPQ p a -> ShowS
[DEPQ p a] -> ShowS
DEPQ p a -> String
(Int -> DEPQ p a -> ShowS)
-> (DEPQ p a -> String) -> ([DEPQ p a] -> ShowS) -> Show (DEPQ p a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall p a. (Show p, Show a) => Int -> DEPQ p a -> ShowS
forall p a. (Show p, Show a) => [DEPQ p a] -> ShowS
forall p a. (Show p, Show a) => DEPQ p a -> String
showList :: [DEPQ p a] -> ShowS
$cshowList :: forall p a. (Show p, Show a) => [DEPQ p a] -> ShowS
show :: DEPQ p a -> String
$cshow :: forall p a. (Show p, Show a) => DEPQ p a -> String
showsPrec :: Int -> DEPQ p a -> ShowS
$cshowsPrec :: forall p a. (Show p, Show a) => Int -> DEPQ p a -> ShowS
Show)
instance Foldable (DEPQ p) where
foldr :: (a -> b -> b) -> b -> DEPQ p a -> b
foldr a -> b -> b
f b
z (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
_) = (a -> b -> b) -> b -> IntPSQ p a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
z IntPSQ p a
mi
instance (NFData p, NFData a) => NFData (DEPQ p a) where
rnf :: DEPQ p a -> ()
rnf (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = IntPSQ p a -> ()
forall a. NFData a => a -> ()
rnf IntPSQ p a
mi () -> () -> ()
`seq` IntPSQ (Down p) a -> ()
forall a. NFData a => a -> ()
rnf IntPSQ (Down p) a
ma
instance (Ord p, Arbitrary p, Arbitrary a) => Arbitrary (DEPQ p a) where
arbitrary :: Gen (DEPQ p a)
arbitrary = [(Int, p, a)] -> DEPQ p a
forall (t :: * -> *) p a.
(Foldable t, Ord p) =>
t (Int, p, a) -> DEPQ p a
fromList ([(Int, p, a)] -> DEPQ p a) -> Gen [(Int, p, a)] -> Gen (DEPQ p a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Gen [(Int, p, a)]
forall a. Arbitrary a => Gen a
arbitrary :: Gen [(Int, p, a)])
shrink :: DEPQ p a -> [DEPQ p a]
shrink DEPQ p a
depq = ([(Int, p, a)] -> DEPQ p a) -> [[(Int, p, a)]] -> [DEPQ p a]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, p, a)] -> DEPQ p a
forall (t :: * -> *) p a.
(Foldable t, Ord p) =>
t (Int, p, a) -> DEPQ p a
fromList ([[(Int, p, a)]] -> [DEPQ p a]) -> [[(Int, p, a)]] -> [DEPQ p a]
forall a b. (a -> b) -> a -> b
$ [(Int, p, a)] -> [[(Int, p, a)]]
forall a. Arbitrary a => a -> [a]
shrink ([(Int, p, a)] -> [[(Int, p, a)]])
-> [(Int, p, a)] -> [[(Int, p, a)]]
forall a b. (a -> b) -> a -> b
$ DEPQ p a -> [(Int, p, a)]
forall p v. DEPQ p v -> [(Int, p, v)]
toList DEPQ p a
depq
insert :: (Ord p) =>
Int
-> p
-> a
-> DEPQ p a -> DEPQ p a
insert :: Int -> p -> a -> DEPQ p a -> DEPQ p a
insert Int
k p
p a
v (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma'
where
mi' :: IntPSQ p a
mi' = Int -> p -> a -> IntPSQ p a -> IntPSQ p a
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
P.insert Int
k p
p a
v IntPSQ p a
mi
ma' :: IntPSQ (Down p) a
ma' = Int -> Down p -> a -> IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
P.insert Int
k (p -> Down p
forall a. a -> Down a
Down p
p) a
v IntPSQ (Down p) a
ma
{-# INLINE insert #-}
empty :: DEPQ p a
empty :: DEPQ p a
empty = IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
forall p v. IntPSQ p v
P.empty IntPSQ (Down p) a
forall p v. IntPSQ p v
P.empty
size :: DEPQ p a -> Int
size :: DEPQ p a -> Int
size (DEPQ IntPSQ p a
p IntPSQ (Down p) a
_) = IntPSQ p a -> Int
forall p v. IntPSQ p v -> Int
P.size IntPSQ p a
p
lookup :: Int
-> DEPQ p v
-> Maybe (p, v)
lookup :: Int -> DEPQ p v -> Maybe (p, v)
lookup Int
k (DEPQ IntPSQ p v
p IntPSQ (Down p) v
_) = Int -> IntPSQ p v -> Maybe (p, v)
forall p v. Int -> IntPSQ p v -> Maybe (p, v)
P.lookup Int
k IntPSQ p v
p
{-# inline lookup #-}
fromList :: (Foldable t, Ord p) =>
t (Int, p, a)
-> DEPQ p a
fromList :: t (Int, p, a) -> DEPQ p a
fromList = (DEPQ p a -> (Int, p, a) -> DEPQ p a)
-> DEPQ p a -> t (Int, p, a) -> DEPQ p a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl DEPQ p a -> (Int, p, a) -> DEPQ p a
forall p a. Ord p => DEPQ p a -> (Int, p, a) -> DEPQ p a
insf DEPQ p a
forall p a. DEPQ p a
empty where
insf :: DEPQ p a -> (Int, p, a) -> DEPQ p a
insf DEPQ p a
acc (Int
k,p
p,a
v) = Int -> p -> a -> DEPQ p a -> DEPQ p a
forall p a. Ord p => Int -> p -> a -> DEPQ p a -> DEPQ p a
insert Int
k p
p a
v DEPQ p a
acc
{-# inline fromList #-}
toList :: DEPQ p v -> [(Int, p, v)]
toList :: DEPQ p v -> [(Int, p, v)]
toList (DEPQ IntPSQ p v
p IntPSQ (Down p) v
_) = IntPSQ p v -> [(Int, p, v)]
forall p v. IntPSQ p v -> [(Int, p, v)]
P.toList IntPSQ p v
p
{-# inline toList #-}
null :: DEPQ p v -> Bool
null :: DEPQ p v -> Bool
null (DEPQ IntPSQ p v
mi IntPSQ (Down p) v
ma) = IntPSQ p v -> Bool
forall p v. IntPSQ p v -> Bool
P.null IntPSQ p v
mi Bool -> Bool -> Bool
&& IntPSQ (Down p) v -> Bool
forall p v. IntPSQ p v -> Bool
P.null IntPSQ (Down p) v
ma
valid :: (Ord p) => DEPQ p v -> Bool
valid :: DEPQ p v -> Bool
valid (DEPQ IntPSQ p v
mi IntPSQ (Down p) v
ma) = IntPSQ p v -> Bool
forall p v. Ord p => IntPSQ p v -> Bool
P.valid IntPSQ p v
mi Bool -> Bool -> Bool
&& IntPSQ (Down p) v -> Bool
forall p v. Ord p => IntPSQ p v -> Bool
P.valid IntPSQ (Down p) v
ma
delete :: Ord p => Int
-> DEPQ p a -> DEPQ p a
delete :: Int -> DEPQ p a -> DEPQ p a
delete Int
k (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma'
where
mi' :: IntPSQ p a
mi' = Int -> IntPSQ p a -> IntPSQ p a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
k IntPSQ p a
mi
ma' :: IntPSQ (Down p) a
ma' = Int -> IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
k IntPSQ (Down p) a
ma
deleteMin :: Ord p => DEPQ p a -> DEPQ p a
deleteMin :: DEPQ p a -> DEPQ p a
deleteMin de :: DEPQ p a
de@(DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = case IntPSQ p a -> Maybe (Int, p, a)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ p a
mi of
Maybe (Int, p, a)
Nothing -> DEPQ p a
de
Just (Int
imin, p
_, a
_) -> IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma' where
mi' :: IntPSQ p a
mi' = IntPSQ p a -> IntPSQ p a
forall p v. Ord p => IntPSQ p v -> IntPSQ p v
P.deleteMin IntPSQ p a
mi
ma' :: IntPSQ (Down p) a
ma' = Int -> IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
imin IntPSQ (Down p) a
ma
{-# INLINE deleteMin #-}
deleteMax :: Ord p => DEPQ p a -> DEPQ p a
deleteMax :: DEPQ p a -> DEPQ p a
deleteMax de :: DEPQ p a
de@(DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = case IntPSQ (Down p) a -> Maybe (Int, Down p, a)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ (Down p) a
ma of
Maybe (Int, Down p, a)
Nothing -> DEPQ p a
de
Just (Int
imax, Down p
_, a
_) -> IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma' where
ma' :: IntPSQ (Down p) a
ma' = IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => IntPSQ p v -> IntPSQ p v
P.deleteMin IntPSQ (Down p) a
ma
mi' :: IntPSQ p a
mi' = Int -> IntPSQ p a -> IntPSQ p a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
imax IntPSQ p a
mi
{-# INLINE deleteMax #-}
findMin :: Ord p => DEPQ p v -> Maybe (Int, p, v)
findMin :: DEPQ p v -> Maybe (Int, p, v)
findMin (DEPQ IntPSQ p v
mi IntPSQ (Down p) v
_) = IntPSQ p v -> Maybe (Int, p, v)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ p v
mi
{-# inline findMin #-}
findMax :: Ord p => DEPQ p v -> Maybe (Int, p, v)
findMax :: DEPQ p v -> Maybe (Int, p, v)
findMax (DEPQ IntPSQ p v
_ IntPSQ (Down p) v
ma) = (Int, Down p, v) -> (Int, p, v)
forall a b c. (a, Down b, c) -> (a, b, c)
f ((Int, Down p, v) -> (Int, p, v))
-> Maybe (Int, Down p, v) -> Maybe (Int, p, v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntPSQ (Down p) v -> Maybe (Int, Down p, v)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ (Down p) v
ma
where
f :: (a, Down b, c) -> (a, b, c)
f (a
i, Down b
p, c
v) = (a
i, b
p, c
v)
{-# inline findMax #-}
popMin :: Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMin :: DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMin DEPQ p v
q = do
(Int, p, v)
x <- DEPQ p v -> Maybe (Int, p, v)
forall p v. Ord p => DEPQ p v -> Maybe (Int, p, v)
findMin DEPQ p v
q
let q' :: DEPQ p v
q' = DEPQ p v -> DEPQ p v
forall p a. Ord p => DEPQ p a -> DEPQ p a
deleteMin DEPQ p v
q
((Int, p, v), DEPQ p v) -> Maybe ((Int, p, v), DEPQ p v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, p, v)
x, DEPQ p v
q')
popMax :: Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMax :: DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMax DEPQ p v
q = do
(Int, p, v)
x <- DEPQ p v -> Maybe (Int, p, v)
forall p v. Ord p => DEPQ p v -> Maybe (Int, p, v)
findMax DEPQ p v
q
let q' :: DEPQ p v
q' = DEPQ p v -> DEPQ p v
forall p a. Ord p => DEPQ p a -> DEPQ p a
deleteMax DEPQ p v
q
((Int, p, v), DEPQ p v) -> Maybe ((Int, p, v), DEPQ p v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, p, v)
x, DEPQ p v
q')
topK :: Ord p => Int -> DEPQ p v -> S.Seq (Int, p, v)
topK :: Int -> DEPQ p v -> Seq (Int, p, v)
topK = (DEPQ p v -> Maybe ((Int, p, v), DEPQ p v))
-> Int -> DEPQ p v -> Seq (Int, p, v)
forall q a. (q -> Maybe (a, q)) -> Int -> q -> Seq a
popK DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
forall p v. Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMax
bottomK :: Ord p => Int -> DEPQ p v -> S.Seq (Int, p, v)
bottomK :: Int -> DEPQ p v -> Seq (Int, p, v)
bottomK = (DEPQ p v -> Maybe ((Int, p, v), DEPQ p v))
-> Int -> DEPQ p v -> Seq (Int, p, v)
forall q a. (q -> Maybe (a, q)) -> Int -> q -> Seq a
popK DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
forall p v. Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMin
popK :: (q -> Maybe (a, q))
-> Int
-> q
-> S.Seq a
popK :: (q -> Maybe (a, q)) -> Int -> q -> Seq a
popK q -> Maybe (a, q)
pop Int
kk q
qq = Seq a -> Maybe (Seq a) -> Seq a
forall a. a -> Maybe a -> a
fromMaybe Seq a
forall a. Seq a
S.empty (Maybe (Seq a) -> Seq a) -> Maybe (Seq a) -> Seq a
forall a b. (a -> b) -> a -> b
$ q -> Int -> Seq a -> Maybe (Seq a)
forall t. (Eq t, Num t) => q -> t -> Seq a -> Maybe (Seq a)
go q
qq Int
kk Seq a
forall a. Seq a
S.empty where
go :: q -> t -> Seq a -> Maybe (Seq a)
go q
_ t
0 Seq a
acc = Seq a -> Maybe (Seq a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Seq a
acc
go q
q t
k Seq a
acc = do
(a
x, q
q') <- q -> Maybe (a, q)
pop q
q
q -> t -> Seq a -> Maybe (Seq a)
go q
q' (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Seq a
acc Seq a -> a -> Seq a
forall a. Seq a -> a -> Seq a
S.|> a
x)