{-# language CPP #-}
{-# OPTIONS_GHC -Wno-unused-imports #-}
module Data.RPTree.Internal.MedianHeap (
  MedianHeap,
  insert, fromList,
  median) where

import Data.Bifunctor (Bifunctor(..))
import Data.Ord (Down(..))

import qualified Data.Heap as H (Heap, Entry(..), viewMin, insert, singleton)

type MinHeap p a = H.Heap (H.Entry p a)
type MaxHeap p a = H.Heap (H.Entry (Down p) a)
-- | The items @a@ in a 'MedianHeap' have type @p@
data MedianHeap p a = MedianHeap (MaxHeap p a) (MinHeap p a) deriving (MedianHeap p a -> MedianHeap p a -> Bool
(MedianHeap p a -> MedianHeap p a -> Bool)
-> (MedianHeap p a -> MedianHeap p a -> Bool)
-> Eq (MedianHeap p a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall p a. MedianHeap p a -> MedianHeap p a -> Bool
/= :: MedianHeap p a -> MedianHeap p a -> Bool
$c/= :: forall p a. MedianHeap p a -> MedianHeap p a -> Bool
== :: MedianHeap p a -> MedianHeap p a -> Bool
$c== :: forall p a. MedianHeap p a -> MedianHeap p a -> Bool
Eq, Int -> MedianHeap p a -> ShowS
[MedianHeap p a] -> ShowS
MedianHeap p a -> String
(Int -> MedianHeap p a -> ShowS)
-> (MedianHeap p a -> String)
-> ([MedianHeap p a] -> ShowS)
-> Show (MedianHeap p a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall p a. (Show p, Show a) => Int -> MedianHeap p a -> ShowS
forall p a. (Show p, Show a) => [MedianHeap p a] -> ShowS
forall p a. (Show p, Show a) => MedianHeap p a -> String
showList :: [MedianHeap p a] -> ShowS
$cshowList :: forall p a. (Show p, Show a) => [MedianHeap p a] -> ShowS
show :: MedianHeap p a -> String
$cshow :: forall p a. (Show p, Show a) => MedianHeap p a -> String
showsPrec :: Int -> MedianHeap p a -> ShowS
$cshowsPrec :: forall p a. (Show p, Show a) => Int -> MedianHeap p a -> ShowS
Show)
instance Semigroup (MedianHeap p a) where
  MedianHeap MaxHeap p a
l1 MinHeap p a
r1 <> :: MedianHeap p a -> MedianHeap p a -> MedianHeap p a
<> MedianHeap MaxHeap p a
l2 MinHeap p a
r2 = MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap (MaxHeap p a
l1 MaxHeap p a -> MaxHeap p a -> MaxHeap p a
forall a. Semigroup a => a -> a -> a
<> MaxHeap p a
l2) (MinHeap p a
r1 MinHeap p a -> MinHeap p a -> MinHeap p a
forall a. Semigroup a => a -> a -> a
<> MinHeap p a
r2)
instance Monoid (MedianHeap p a) where
  mempty :: MedianHeap p a
mempty = MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap MaxHeap p a
forall a. Monoid a => a
mempty MinHeap p a
forall a. Monoid a => a
mempty
insertMax :: Ord p => p -> a -> MaxHeap p a -> MaxHeap p a
insertMax :: p -> a -> MaxHeap p a -> MaxHeap p a
insertMax p
p a
x = Entry (Down p) a -> MaxHeap p a -> MaxHeap p a
forall a. Ord a => a -> Heap a -> Heap a
H.insert (Down p -> a -> Entry (Down p) a
forall p a. p -> a -> Entry p a
H.Entry (p -> Down p
forall a. a -> Down a
Down p
p) a
x)
insertMin :: Ord p => p -> a -> MinHeap p a -> MinHeap p a
insertMin :: p -> a -> MinHeap p a -> MinHeap p a
insertMin p
p a
x = Entry p a -> MinHeap p a -> MinHeap p a
forall a. Ord a => a -> Heap a -> Heap a
H.insert (p -> a -> Entry p a
forall p a. p -> a -> Entry p a
H.Entry p
p a
x)

fromList :: (Foldable t, Fractional p, Ord p) => t (p, a) -> MedianHeap p a
fromList :: t (p, a) -> MedianHeap p a
fromList = ((p, a) -> MedianHeap p a -> MedianHeap p a)
-> MedianHeap p a -> t (p, a) -> MedianHeap p a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((p -> a -> MedianHeap p a -> MedianHeap p a)
-> (p, a) -> MedianHeap p a -> MedianHeap p a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry p -> a -> MedianHeap p a -> MedianHeap p a
forall p a.
(Fractional p, Ord p) =>
p -> a -> MedianHeap p a -> MedianHeap p a
insert) MedianHeap p a
forall a. Monoid a => a
mempty

-- | Insert a weighted entry in the heap
insert :: (Fractional p, Ord p) =>
          p -- ^ weight
       -> a
       -> MedianHeap p a
       -> MedianHeap p a
insert :: p -> a -> MedianHeap p a -> MedianHeap p a
insert p
p a
x heap :: MedianHeap p a
heap@(MedianHeap MaxHeap p a
ll MinHeap p a
rr) =
  case MedianHeap p a -> Maybe p
forall p a. Fractional p => MedianHeap p a -> Maybe p
median MedianHeap p a
heap of
    Maybe p
Nothing -> MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap (p -> a -> MaxHeap p a -> MaxHeap p a
forall p a. Ord p => p -> a -> MaxHeap p a -> MaxHeap p a
insertMax p
p a
x MaxHeap p a
forall a. Monoid a => a
mempty) MinHeap p a
rr
    Just p
q ->
      if p
q p -> p -> Bool
forall a. Ord a => a -> a -> Bool
> p
p
      then MedianHeap p a -> MedianHeap p a
forall p a. Ord p => MedianHeap p a -> MedianHeap p a
balance (MedianHeap p a -> MedianHeap p a)
-> MedianHeap p a -> MedianHeap p a
forall a b. (a -> b) -> a -> b
$ MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap MaxHeap p a
ll (p -> a -> MinHeap p a -> MinHeap p a
forall p a. Ord p => p -> a -> MinHeap p a -> MinHeap p a
insertMin p
p a
x MinHeap p a
rr)
      else MedianHeap p a -> MedianHeap p a
forall p a. Ord p => MedianHeap p a -> MedianHeap p a
balance (MedianHeap p a -> MedianHeap p a)
-> MedianHeap p a -> MedianHeap p a
forall a b. (a -> b) -> a -> b
$ MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap (p -> a -> MaxHeap p a -> MaxHeap p a
forall p a. Ord p => p -> a -> MaxHeap p a -> MaxHeap p a
insertMax p
p a
x MaxHeap p a
ll) MinHeap p a
rr

balance :: Ord p => MedianHeap p a -> MedianHeap p a
balance :: MedianHeap p a -> MedianHeap p a
balance (MedianHeap MaxHeap p a
ll MinHeap p a
rr)
  | MaxHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MaxHeap p a
ll Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== MinHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MinHeap p a
rr = MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap MaxHeap p a
ll MinHeap p a
rr
  | MaxHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MaxHeap p a
ll Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> MinHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MinHeap p a
rr =
    let (H.Entry (Down p
p) a
x, MaxHeap p a
xs) = MaxHeap p a -> (Entry (Down p) a, MaxHeap p a)
forall a. Heap a -> (a, Heap a)
deconstruct MaxHeap p a
ll
     in MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap MaxHeap p a
xs (p -> a -> MinHeap p a -> MinHeap p a
forall p a. Ord p => p -> a -> MinHeap p a -> MinHeap p a
insertMin p
p a
x MinHeap p a
rr)
  | Bool
otherwise =
    let (H.Entry p
p a
x, MinHeap p a
xs) = MinHeap p a -> (Entry p a, MinHeap p a)
forall a. Heap a -> (a, Heap a)
deconstruct MinHeap p a
rr
     in MaxHeap p a -> MinHeap p a -> MedianHeap p a
forall p a. MaxHeap p a -> MinHeap p a -> MedianHeap p a
MedianHeap (p -> a -> MaxHeap p a -> MaxHeap p a
forall p a. Ord p => p -> a -> MaxHeap p a -> MaxHeap p a
insertMax p
p a
x MaxHeap p a
ll) MinHeap p a
xs
  where
    deconstruct :: Heap a -> (a, Heap a)
deconstruct Heap a
heap = case Heap a -> Maybe (a, Heap a)
forall a. Heap a -> Maybe (a, Heap a)
H.viewMin Heap a
heap of
      Maybe (a, Heap a)
Nothing -> String -> (a, Heap a)
forall a. HasCallStack => String -> a
error String
"cannot view empty heap"
      Just (a
x, Heap a
xs) -> (a
x, Heap a
xs)

-- | Compute the median weight
median :: Fractional p => MedianHeap p a -> Maybe p
median :: MedianHeap p a -> Maybe p
median (MedianHeap MaxHeap p a
lesser MinHeap p a
greater)
  | MaxHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MaxHeap p a
lesser Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> MinHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MinHeap p a
greater = Entry p a -> p
forall p a. Entry p a -> p
H.priority (Entry p a -> p)
-> (Entry (Down p) a -> Entry p a) -> Entry (Down p) a -> p
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry (Down p) a -> Entry p a
forall b c. Entry (Down b) c -> Entry b c
getHD (Entry (Down p) a -> p) -> Maybe (Entry (Down p) a) -> Maybe p
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MaxHeap p a -> Maybe (Entry (Down p) a)
forall b. Heap b -> Maybe b
viewHead MaxHeap p a
lesser
  | MaxHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MaxHeap p a
lesser Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MinHeap p a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MinHeap p a
greater = Entry p a -> p
forall p a. Entry p a -> p
H.priority (Entry p a -> p) -> Maybe (Entry p a) -> Maybe p
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MinHeap p a -> Maybe (Entry p a)
forall b. Heap b -> Maybe b
viewHead MinHeap p a
greater
  | Bool
otherwise = do
      Entry p a
leftHead <- Entry (Down p) a -> Entry p a
forall b c. Entry (Down b) c -> Entry b c
getHD (Entry (Down p) a -> Entry p a)
-> Maybe (Entry (Down p) a) -> Maybe (Entry p a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MaxHeap p a -> Maybe (Entry (Down p) a)
forall b. Heap b -> Maybe b
viewHead MaxHeap p a
lesser
      Entry p a
rightHead <- MinHeap p a -> Maybe (Entry p a)
forall b. Heap b -> Maybe b
viewHead MinHeap p a
greater
      p -> Maybe p
forall (m :: * -> *) a. Monad m => a -> m a
return (p -> Maybe p) -> p -> Maybe p
forall a b. (a -> b) -> a -> b
$ (Entry p a -> p
forall p a. Entry p a -> p
H.priority Entry p a
leftHead p -> p -> p
forall a. Num a => a -> a -> a
+ Entry p a -> p
forall p a. Entry p a -> p
H.priority Entry p a
rightHead) p -> p -> p
forall a. Fractional a => a -> a -> a
/ p
2

getHD :: H.Entry (Down b) c -> H.Entry b c
getHD :: Entry (Down b) c -> Entry b c
getHD = (Down b -> b) -> Entry (Down b) c -> Entry b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Down b -> b
forall a. Down a -> a
getDown

#if MIN_VERSION_base(4,14,0)
#else
getDown :: Down a -> a
getDown (Down x) = x
#endif

viewHead :: H.Heap b -> Maybe b
viewHead :: Heap b -> Maybe b
viewHead Heap b
h = (b, Heap b) -> b
forall a b. (a, b) -> a
fst ((b, Heap b) -> b) -> Maybe (b, Heap b) -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Heap b -> Maybe (b, Heap b)
forall a. Heap a -> Maybe (a, Heap a)
H.viewMin Heap b
h