module Data.IntHeap.MaxTopK where

import qualified Data.List as L

import Data.IntHeap

data MaxTopK = MaxTopK !IntHeap !Int !Int !IntHeap
  deriving (MaxTopK -> MaxTopK -> Bool
(MaxTopK -> MaxTopK -> Bool)
-> (MaxTopK -> MaxTopK -> Bool) -> Eq MaxTopK
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MaxTopK -> MaxTopK -> Bool
== :: MaxTopK -> MaxTopK -> Bool
$c/= :: MaxTopK -> MaxTopK -> Bool
/= :: MaxTopK -> MaxTopK -> Bool
Eq)

instance Show MaxTopK where
  show :: MaxTopK -> String
show (MaxTopK IntHeap
minHeapL Int
minL Int
maxR IntHeap
maxHeapR) =
    ([Int], [Int]) -> String
forall a. Show a => a -> String
show
      ( [Int] -> [Int]
forall a. [a] -> [a]
reverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Int
minL Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: IntHeap -> [Int]
toAscListIH IntHeap
minHeapL
      , Int
maxR Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: IntHeap -> [Int]
toDescListIH IntHeap
maxHeapR
      )

{- | /O(1)/

>>> viewMaxTopK (buildMaxTopK 3 [0..9])
7
-}
viewMaxTopK :: MaxTopK -> Int
viewMaxTopK :: MaxTopK -> Int
viewMaxTopK (MaxTopK IntHeap
_ Int
minL Int
_ IntHeap
_) = Int
minL

{- | /O(k)/

>>> viewMaxTopKList (buildMaxTopK 3 [0..9])
[9,8,7]
-}
viewMaxTopKList :: MaxTopK -> [Int]
viewMaxTopKList :: MaxTopK -> [Int]
viewMaxTopKList (MaxTopK IntHeap
minHeapL Int
minL Int
_ IntHeap
_) = IntHeap -> [Int]
toDescListIH IntHeap
minHeapL [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
minL]

newMaxTopK :: Int -> MaxTopK
newMaxTopK :: Int -> MaxTopK
newMaxTopK Int
k = IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK (Int -> Int -> IntHeap
replicateIH (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
forall a. Bounded a => a
minBound) Int
forall a. Bounded a => a
minBound Int
forall a. Bounded a => a
minBound IntHeap
emptyIH

{- | /O(n log n)/

>>> buildMaxTopK 4 [0..4]
([4,3,2,1],[0])
>>> buildMaxTopK 4 [0..3]
([3,2,1,0],[-9223372036854775808])
>>> buildMaxTopK 4 [0..2]
([2,1,0,-9223372036854775808],[-9223372036854775808])
-}
buildMaxTopK :: Int -> [Int] -> MaxTopK
buildMaxTopK :: Int -> [Int] -> MaxTopK
buildMaxTopK Int
k [Int]
xs =
  Int -> [Int] -> MaxTopK
buildMaxTopKFromDescList Int
k ([Int] -> MaxTopK) -> [Int] -> MaxTopK
forall a b. (a -> b) -> a -> b
$
    (Int -> Int -> Ordering) -> [Int] -> [Int]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy ((Int -> Int -> Ordering) -> Int -> Int -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare) [Int]
xs

{- | /O(n)/

>>> buildMaxTopKFromDescList 3 [2,2,1,1,0,0]
([2,2,1],[1,0,0])
-}
buildMaxTopKFromDescList :: Int -> [Int] -> MaxTopK
buildMaxTopKFromDescList :: Int -> [Int] -> MaxTopK
buildMaxTopKFromDescList Int
k [Int]
xs = case Int -> [Int] -> ([Int], [Int])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Int]
xs of
  ([Int]
ls, Int
minL : Int
maxR : [Int]
rs) -> IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK ([Int] -> IntHeap
fromDescListIH [Int]
ls) Int
minL Int
maxR ([Int] -> IntHeap
fromDescListIH [Int]
rs)
  ([Int]
ls, [Int
minL]) -> IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK ([Int] -> IntHeap
fromDescListIH [Int]
ls) Int
minL Int
forall a. Bounded a => a
minBound IntHeap
emptyIH
  ([Int]
ls, []) ->
    IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK
      ([Int] -> IntHeap
fromDescListIH ([Int] -> IntHeap) -> ([Int] -> [Int]) -> [Int] -> IntHeap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([Int] -> IntHeap) -> [Int] -> IntHeap
forall a b. (a -> b) -> a -> b
$ [Int]
ls [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
forall a. Bounded a => a
minBound)
      Int
forall a. Bounded a => a
minBound
      Int
forall a. Bounded a => a
minBound
      IntHeap
emptyIH

{- | /O(min(n,W))/

>>> maxTop4 = buildMaxTopK 4 [0..4]
>>> maxTop4
([4,3,2,1],[0])
>>> insertMaxTopK 999 maxTop4
(Just 1,([999,4,3,2],[1,0]))
>>> insertMaxTopK (-999) maxTop4
(Nothing,([4,3,2,1],[0,-999]))
>>> insertMaxTopK 1 maxTop4
(Nothing,([4,3,2,1],[1,0]))
-}
insertMaxTopK ::
  Int ->
  MaxTopK ->
  -- | Maybe deleted
  (Maybe Int, MaxTopK)
insertMaxTopK :: Int -> MaxTopK -> (Maybe Int, MaxTopK)
insertMaxTopK Int
x (MaxTopK IntHeap
minHeapL Int
minL Int
maxR IntHeap
maxHeapR)
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
minL =
    ( Maybe Int
forall a. Maybe a
Nothing
    , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK IntHeap
minHeapL Int
minL (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
x Int
maxR) (Int -> IntHeap -> IntHeap
insertIH (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
x Int
maxR) IntHeap
maxHeapR)
    )
  | Bool
otherwise =
    case IntHeap -> (Int, IntHeap)
deleteFindMinIH (Int -> IntHeap -> IntHeap
insertIH Int
x IntHeap
minHeapL) of
      (Int
minL', IntHeap
minHeapL') ->
        ( Int -> Maybe Int
forall a. a -> Maybe a
Just Int
minL
        , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK IntHeap
minHeapL' Int
minL' Int
minL (Int -> IntHeap -> IntHeap
insertIH Int
maxR IntHeap
maxHeapR)
        )

{- | /O(min(n,W))/

>>> maxTop3 = buildMaxTopK 3 [0,1,1,2,2]
>>> maxTop3
([2,2,1],[1,0])
>>> deleteMaxTopK 2 maxTop3
(Just 1,([2,1,1],[0]))
>>> deleteMaxTopK 1 maxTop3
(Nothing,([2,2,1],[0]))
>>> deleteMaxTopK 0 maxTop3
(Nothing,([2,2,1],[1]))
>>> deleteMaxTopK 999 maxTop3
(Nothing,([2,2,1],[1,0]))
>>> deleteMaxTopK (-999) maxTop3
(Nothing,([2,2,1],[1,0]))
>>> deleteMaxTopK 2 $ buildMaxTopK 3 [2,2,2,2,2]
(Nothing,([2,2,2],[2]))
>>> deleteMaxTopK 0 $ buildMaxTopK 1 [0]
(Just (-9223372036854775808),([-9223372036854775808],[-9223372036854775808]))
-}
deleteMaxTopK ::
  Int ->
  MaxTopK ->
  -- | Maybe inserted
  (Maybe Int, MaxTopK)
deleteMaxTopK :: Int -> MaxTopK -> (Maybe Int, MaxTopK)
deleteMaxTopK Int
x (MaxTopK IntHeap
minHeapL Int
minL Int
maxR IntHeap
maxHeapR)
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxR =
    ( Maybe Int
forall a. Maybe a
Nothing
    , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK IntHeap
minHeapL Int
minL Int
maxR (IntHeap -> MaxTopK) -> IntHeap -> MaxTopK
forall a b. (a -> b) -> a -> b
$ Int -> IntHeap -> IntHeap
deleteIH Int
x IntHeap
maxHeapR
    )
  | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxR = case IntHeap -> (Int, IntHeap)
deleteFindMax IntHeap
maxHeapR of
    (Int
maxR', IntHeap
maxHeapR') ->
      ( Maybe Int
forall a. Maybe a
Nothing
      , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK IntHeap
minHeapL Int
minL Int
maxR' IntHeap
maxHeapR'
      )
  | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
minL = case IntHeap -> (Int, IntHeap)
deleteFindMax IntHeap
maxHeapR of
    (Int
maxR', IntHeap
maxHeapR') ->
      ( Int -> Maybe Int
forall a. a -> Maybe a
Just Int
maxR
      , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK IntHeap
minHeapL Int
maxR Int
maxR' IntHeap
maxHeapR'
      )
  | Int -> IntHeap -> Bool
notMemberIH Int
x IntHeap
minHeapL =
    ( Maybe Int
forall a. Maybe a
Nothing
    , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK IntHeap
minHeapL Int
minL Int
maxR IntHeap
maxHeapR
    )
  | Bool
otherwise = case IntHeap -> (Int, IntHeap)
deleteFindMax IntHeap
maxHeapR of
    (Int
maxR', IntHeap
maxHeapR') ->
      ( Int -> Maybe Int
forall a. a -> Maybe a
Just Int
maxR
      , IntHeap -> Int -> Int -> IntHeap -> MaxTopK
MaxTopK (Int -> IntHeap -> IntHeap
insertIH Int
minL (IntHeap -> IntHeap) -> IntHeap -> IntHeap
forall a b. (a -> b) -> a -> b
$ Int -> IntHeap -> IntHeap
deleteIH Int
x IntHeap
minHeapL) Int
maxR Int
maxR' IntHeap
maxHeapR'
      )
  where
    deleteFindMax :: IntHeap -> (Int, IntHeap)
deleteFindMax IntHeap
heap = case IntHeap -> Maybe (Int, IntHeap)
maxViewIH IntHeap
heap of
      Just (Int
mm, IntHeap
heap') -> (Int
mm, IntHeap
heap')
      Maybe (Int, IntHeap)
Nothing -> (Int
forall a. Bounded a => a
minBound, IntHeap
heap)

{- | /O(min(n,W))/

>>> maxTop3 = buildMaxTopK 3 [0,1,1,2,2]
>>> maxTop3
([2,2,1],[1,0])
>>> exchangeMaxTopK 2 999 maxTop3
(Just (999,2),([999,2,1],[1,0]))
>>> exchangeMaxTopK 1 999 maxTop3
(Just (999,1),([999,2,2],[1,0]))
>>> exchangeMaxTopK 0 999 maxTop3
(Just (999,1),([999,2,2],[1,1]))
>>> exchangeMaxTopK 2 1 maxTop3
(Just (1,2),([2,1,1],[1,0]))
>>> exchangeMaxTopK 0 1 maxTop3
(Nothing,([2,2,1],[1,1]))
>>> exchangeMaxTopK 1 1 maxTop3
(Nothing,([2,2,1],[1,0]))
>>> exchangeMaxTopK 999 999999 maxTop3
(Nothing,([2,2,1],[1,0]))
>>> exchangeMaxTopK 0 999 $ buildMaxTopK 1 [0, 0]
(Just (999,0),([999],[0]))
-}
exchangeMaxTopK ::
  -- | old
  Int ->
  -- | new
  Int ->
  MaxTopK ->
  -- | Maybe (inserted, deleted)
  (Maybe (Int, Int), MaxTopK)
exchangeMaxTopK :: Int -> Int -> MaxTopK -> (Maybe (Int, Int), MaxTopK)
exchangeMaxTopK Int
old Int
new maxTopK :: MaxTopK
maxTopK@(MaxTopK IntHeap
minHeapL Int
minL Int
maxR IntHeap
maxHeapR)
  | Int
minL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
old
    , Int
maxR Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
old
    , Int
old Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minL Bool -> Bool -> Bool
|| Int -> IntHeap -> Bool
notMemberIH Int
old IntHeap
minHeapL
    , Int
old Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxR Bool -> Bool -> Bool
|| Int -> IntHeap -> Bool
notMemberIH Int
old IntHeap
maxHeapR =
    (Maybe (Int, Int)
forall a. Maybe a
Nothing, MaxTopK
maxTopK)
  | Int
old Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
new = (Maybe (Int, Int)
forall a. Maybe a
Nothing, MaxTopK
maxTopK)
exchangeMaxTopK Int
old Int
new MaxTopK
maxTopK = case Int -> MaxTopK -> (Maybe Int, MaxTopK)
insertMaxTopK Int
new MaxTopK
maxTopK of
  (Maybe Int
Nothing, MaxTopK
maxTopK') -> case Int -> MaxTopK -> (Maybe Int, MaxTopK)
deleteMaxTopK Int
old MaxTopK
maxTopK' of
    (Maybe Int
Nothing, MaxTopK
maxTopK'') -> (Maybe (Int, Int)
forall a. Maybe a
Nothing, MaxTopK
maxTopK'')
    (Just Int
inserted, MaxTopK
maxTopK'') -> ((Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
inserted, Int
old), MaxTopK
maxTopK'')
  (Just Int
deleted, MaxTopK
maxTopK') -> case Int -> MaxTopK -> (Maybe Int, MaxTopK)
deleteMaxTopK Int
old MaxTopK
maxTopK' of
    (Maybe Int
Nothing, MaxTopK
maxTopK'') -> ((Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
new, Int
deleted), MaxTopK
maxTopK'')
    (Just Int
_inserted, MaxTopK
maxTopK'') -> ((Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
new, Int
old), MaxTopK
maxTopK'')