module Data.Queue.SoftHeap (SoftHeap, empty', singleton', fromList') where
import Data.Queue.Class
import Data.Queue.Numeric
import Data.Sequence (Seq, viewl, (<|), ViewL(..), ViewR(..))
import qualified Data.Sequence as Seq
import qualified Data.Foldable as Fold
import Data.Ord
import Data.Ratio
import Data.Maybe
import Control.Monad
import Control.Arrow(second)
import Control.Monad.Instances
import Data.Tree hiding (subForest)
data SNode e = SN {ckey :: e, rk, targetSize :: !Int, list :: [e], left, right :: Maybe (SNode e)} deriving (Show)
data SHead e = SHead {sHeap :: !(SNode e), myIx, sufMin :: !Int} deriving (Show)
type SHeapList e = Seq (SHead e)
data SoftHeap e = SQ {elts, rank, rConst :: !Int, heads :: SHeapList e}
toTree :: SNode e -> Tree (e, Int, [e])
toTree SN{..} = Node (ckey, rk, list) [toTree n | Just n <- [left, right]]
toForest :: SoftHeap e -> Forest (e, Int, [e])
toForest = map (toTree . sHeap) . Fold.toList . heads
drawHeap :: Show e => SoftHeap e -> String
drawHeap = drawForest . map (fmap show) . toForest
defaultRank :: Int
defaultRank = 12
fromEpsilon :: RealFrac b => b -> Int
fromEpsilon x
= let (a, b) = asFraction (toPrecision x) in ceilLog (fromIntegral b) intLog (fromIntegral a) + 5
where precision = fromRational (1 % 100000)
toPrecision x = x `approxRational` (precision `min` (x / 2))
asFraction = liftM2 (,) numerator denominator
empty' :: (Ord e, RealFrac b) => b -> SoftHeap e
empty' epsilon = SQ 0 0 (fromEpsilon epsilon) Seq.empty
singleton' :: (Ord e, RealFrac b) => b -> e -> SoftHeap e
singleton' epsilon x = SQ 1 0 (fromEpsilon epsilon) (Seq.singleton (SHead (single x) 0 0))
fromList' :: (Ord e, RealFrac b) => b -> [e] -> SoftHeap e
fromList' epsilon xs = insertAll xs (empty' epsilon)
instance (Ord e) => Queuelike (SoftHeap e) where
type QueueKey (SoftHeap e) = e
empty = SQ 0 0 defaultRank Seq.empty
singleton x = SQ 1 0 defaultRank (Seq.singleton (SHead (single x) 0 0))
merge = meld
extract = deleteMin
size = elts
minWith :: Ord b => (a -> b) -> a -> a -> a
minWith f x y | f x <= f y = x
| otherwise = y
orderPairWith :: Ord b => (a -> b) -> (a, a) -> (a, a)
orderPairWith f (x, y) | f x <= f y = (x, y)
| otherwise = (y, x)
isLeaf :: SNode e -> Bool
isLeaf SN{..} = not (isJust left || isJust right)
sift :: Ord e => SNode e -> SNode e
sift x@SN{left, right, targetSize}
| length (list x) >= targetSize || isLeaf x
= x
| (l', r') <- swapKids left right
= let ckey' = case list x of
[] -> ckey l'
_ -> ckey x `max` ckey l'
in ckey' `seq` sift x{ckey = ckey', list = list x ++ list l',
left = if isLeaf l' then Nothing else Just (sift l'{list = []}), right = r'}
where swapKids (Just l) (Just r) = let (l', r') = orderPairWith ckey (l, r) in (l', Just r')
swapKids (Just l) r = (l, r)
swapKids _ (Just r) = (r, Nothing)
combine :: Ord e => Int -> SNode e -> SNode e -> SNode e
combine r x@SN{rk,targetSize} y = sift $ SN undefined (rk + 1) (if rk < r then 1 else (3 * targetSize + 1) `quot` 2) [] (Just x) (Just y)
consHead :: Ord e => SNode e -> SHeapList e -> SHeapList e
consHead x hs = case viewl hs of
EmptyL -> SHead x 0 0 <| hs
(SHead{sufMin} :< _) | n <- Seq.length hs, SHead{sHeap = sMin} <- Seq.index hs (n 1 sufMin)
-> SHead x n (if ckey x <= ckey sMin then n else sufMin) <| hs
single :: Ord e => e -> SNode e
single x = SN x 0 1 [x] Nothing Nothing
rkHead :: SHead e -> Int
rkHead = rk . sHeap
meld :: (Ord e) => SoftHeap e -> SoftHeap e -> SoftHeap e
SQ n1 rk1 r1 p `meld` SQ n2 rk2 r2 q = case viewl (p `mergeRanks` q) of
EmptyL -> SQ 0 0 r Seq.empty
(SHead h _ _ :< hs) | (rk, heads) <- compare (ckey h) (ckey h) `seq` rebuild h hs
-> SQ (n1 + n2) (rk `max` rk1 `max` rk2) r heads
where r = r1 `max` r2
rkMin = rk1 `min` rk2
ps `mergeRanks` qs = case (viewl ps, viewl qs) of
(p :< ps', q :< qs')
| rkHead p <= rkHead q -> p <| mergeRanks ps' qs
| otherwise -> q <| mergeRanks ps qs'
(EmptyL, _) -> qs
(_, EmptyL) -> ps
rebuild q@SN{rk = rk0} qs@(viewl -> SHead{sHeap = q1} :< qs1)
| rk0 == rk q1
= case viewl qs1 of
(SHead{sHeap = q2} :< qs2)
| rk0 == rk q2 -> fmap (q `consHead`) $ rebuild (combine r q1 q2) qs2
_ -> rebuild (combine r q q1) qs1
| otherwise
= fmap (q `consHead`) $ rebuild q1 qs1
rebuild q qs
= (rk q, q `consHead` Seq.empty)
headKey :: SHead e -> e
headKey = ckey . sHeap
fixSufMins :: Ord e => Int -> SHeapList e -> SHeapList e
fixSufMins i sequ@(Seq.splitAt i -> (seqL, seqR))
= Fold.foldr (consHead . sHeap) seqR seqL
deleteMin :: Ord e => SoftHeap e -> Maybe (e, SoftHeap e)
deleteMin q@SQ{elts, rConst = r, heads = heads@(viewl -> SHead{sufMin} :< _)}
| n <- Seq.length heads, t@(headKey -> minKey) <- Seq.index heads (n 1 sufMin)
= Just (minKey, q{elts = elts 1, heads = deleteMin' (n 1 sufMin) t heads})
where deleteMin' :: Ord e => Int -> SHead e -> SHeapList e -> SHeapList e
deleteMin' sufMin t@SHead{sHeap = h@SN{ckey, targetSize, list = _:l'}} heads
| 2 * length l' <= targetSize, not (isLeaf h), h' <- sift h{list = l'}
= case list h' of
[] -> let (lHeads, viewl -> _ :< rHeads) = Seq.splitAt sufMin heads in
Fold.foldr (consHead . sHeap) rHeads lHeads
_ -> Seq.update sufMin t{sHeap = h'} heads
| [] <- l', (lHeads, viewl -> _ :< rHeads) <- Seq.splitAt sufMin heads
= Fold.foldr (consHead . sHeap) rHeads lHeads
| otherwise
= Seq.update sufMin t{sHeap = h{list = l'}} heads
deleteMin _ = Nothing