{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoFieldSelectors #-}
module Data.List.ApplyMerge.IntSet (applyMerge) where
import Control.Monad (guard)
import Data.Function ((&))
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.List (unfoldr)
import Data.List.NonEmpty (NonEmpty, nonEmpty)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Maybe (fromMaybe)
import Data.PQueue.Prio.Min (MinPQueue)
import Data.PQueue.Prio.Min qualified as MinPQueue
data Node a b c = Node
{ forall a b c. Node a b c -> (Int, Int)
position :: (Int, Int),
forall a b c. Node a b c -> c
value :: c,
forall a b c. Node a b c -> NonEmpty a
as :: NonEmpty a,
forall a b c. Node a b c -> NonEmpty b
bs :: NonEmpty b
}
data Frontier a b c = Frontier
{ forall a b c. Frontier a b c -> MinPQueue c (Node a b c)
queue :: MinPQueue c (Node a b c),
forall a b c. Frontier a b c -> IntSet
indexSetA :: IntSet,
forall a b c. Frontier a b c -> IntSet
indexSetB :: IntSet
}
applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge :: forall c a b. Ord c => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge a -> b -> c
f [a]
as [b]
bs = [c] -> Maybe [c] -> [c]
forall a. a -> Maybe a -> a
fromMaybe [] (Maybe [c] -> [c]) -> Maybe [c] -> [c]
forall a b. (a -> b) -> a -> b
$ do
NonEmpty a
as' <- [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [a]
as
NonEmpty b
bs' <- [b] -> Maybe (NonEmpty b)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [b]
bs
[c] -> Maybe [c]
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Frontier a b c -> Maybe (c, Frontier a b c))
-> Frontier a b c -> [c]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ((a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
forall c a b.
Ord c =>
(a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
step a -> b -> c
f) ((a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
initialFrontier a -> b -> c
f NonEmpty a
as' NonEmpty b
bs'))
initialFrontier :: (a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
initialFrontier :: forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
initialFrontier a -> b -> c
f NonEmpty a
as NonEmpty b
bs =
let node :: Node a b c
node = (a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
0, Int
0) NonEmpty a
as NonEmpty b
bs
in Frontier
{ $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = c -> Node a b c -> MinPQueue c (Node a b c)
forall k a. k -> a -> MinPQueue k a
MinPQueue.singleton Node a b c
node.value Node a b c
node,
$sel:indexSetA:Frontier :: IntSet
indexSetA = Int -> IntSet
IntSet.singleton Int
0,
$sel:indexSetB:Frontier :: IntSet
indexSetB = Int -> IntSet
IntSet.singleton Int
0
}
step :: (Ord c) => (a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
step :: forall c a b.
Ord c =>
(a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
step a -> b -> c
f Frontier a b c
frontier = do
(Node a b c
node, Frontier a b c
frontier') <- Frontier a b c -> Maybe (Node a b c, Frontier a b c)
forall c a b.
Ord c =>
Frontier a b c -> Maybe (Node a b c, Frontier a b c)
deleteMinNode Frontier a b c
frontier
let frontier'' :: Frontier a b c
frontier'' =
Frontier a b c
frontier'
Frontier a b c
-> (Frontier a b c -> Frontier a b c) -> Frontier a b c
forall a b. a -> (a -> b) -> b
& (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildA a -> b -> c
f Node a b c
node
Frontier a b c
-> (Frontier a b c -> Frontier a b c) -> Frontier a b c
forall a b. a -> (a -> b) -> b
& (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildB a -> b -> c
f Node a b c
node
(c, Frontier a b c) -> Maybe (c, Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Node a b c
node.value, Frontier a b c
frontier'')
deleteMinNode :: (Ord c) => Frontier a b c -> Maybe (Node a b c, Frontier a b c)
deleteMinNode :: forall c a b.
Ord c =>
Frontier a b c -> Maybe (Node a b c, Frontier a b c)
deleteMinNode Frontier a b c
frontier = do
(Node a b c
node, MinPQueue c (Node a b c)
queue') <- MinPQueue c (Node a b c)
-> Maybe (Node a b c, MinPQueue c (Node a b c))
forall k a. Ord k => MinPQueue k a -> Maybe (a, MinPQueue k a)
MinPQueue.minView Frontier a b c
frontier.queue
let (Int
ia, Int
ib) = Node a b c
node.position
frontier' :: Frontier a b c
frontier' =
Frontier
{ $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = MinPQueue c (Node a b c)
queue',
$sel:indexSetA:Frontier :: IntSet
indexSetA = Int -> IntSet -> IntSet
IntSet.delete Int
ia Frontier a b c
frontier.indexSetA,
$sel:indexSetB:Frontier :: IntSet
indexSetB = Int -> IntSet -> IntSet
IntSet.delete Int
ib Frontier a b c
frontier.indexSetB
}
(Node a b c, Frontier a b c) -> Maybe (Node a b c, Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Node a b c
node, Frontier a b c
frontier')
insertChildA ::
(Ord c) => (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildA :: forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildA a -> b -> c
f (Node (Int
ia, Int
ib) c
_ NonEmpty a
as NonEmpty b
bs) Frontier a b c
frontier = Frontier a b c -> Maybe (Frontier a b c) -> Frontier a b c
forall a. a -> Maybe a -> a
fromMaybe Frontier a b c
frontier (Maybe (Frontier a b c) -> Frontier a b c)
-> Maybe (Frontier a b c) -> Frontier a b c
forall a b. (a -> b) -> a -> b
$ do
Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Int -> IntSet -> Bool
IntSet.member (Int
ia Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Frontier a b c
frontier.indexSetA))
NonEmpty a
as' <- [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NonEmpty.tail NonEmpty a
as)
let childA :: Node a b c
childA = (a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
ia Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ib) NonEmpty a
as' NonEmpty b
bs
Frontier a b c -> Maybe (Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Frontier a b c -> Maybe (Frontier a b c))
-> Frontier a b c -> Maybe (Frontier a b c)
forall a b. (a -> b) -> a -> b
$ Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
Node a b c -> Frontier a b c -> Frontier a b c
insertNode Node a b c
childA Frontier a b c
frontier
insertChildB ::
(Ord c) => (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildB :: forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildB a -> b -> c
f (Node (Int
ia, Int
ib) c
_ NonEmpty a
as NonEmpty b
bs) Frontier a b c
frontier = Frontier a b c -> Maybe (Frontier a b c) -> Frontier a b c
forall a. a -> Maybe a -> a
fromMaybe Frontier a b c
frontier (Maybe (Frontier a b c) -> Frontier a b c)
-> Maybe (Frontier a b c) -> Frontier a b c
forall a b. (a -> b) -> a -> b
$ do
Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Int -> IntSet -> Bool
IntSet.member (Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Frontier a b c
frontier.indexSetB))
NonEmpty b
bs' <- [b] -> Maybe (NonEmpty b)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty (NonEmpty b -> [b]
forall a. NonEmpty a -> [a]
NonEmpty.tail NonEmpty b
bs)
let childB :: Node a b c
childB = (a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
ia, Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) NonEmpty a
as NonEmpty b
bs'
Frontier a b c -> Maybe (Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Frontier a b c -> Maybe (Frontier a b c))
-> Frontier a b c -> Maybe (Frontier a b c)
forall a b. (a -> b) -> a -> b
$ Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
Node a b c -> Frontier a b c -> Frontier a b c
insertNode Node a b c
childB Frontier a b c
frontier
insertNode :: (Ord c) => Node a b c -> Frontier a b c -> Frontier a b c
insertNode :: forall c a b.
Ord c =>
Node a b c -> Frontier a b c -> Frontier a b c
insertNode Node a b c
node Frontier a b c
frontier =
let (Int
ia, Int
ib) = Node a b c
node.position
in Frontier
{ $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = c
-> Node a b c
-> MinPQueue c (Node a b c)
-> MinPQueue c (Node a b c)
forall k a. Ord k => k -> a -> MinPQueue k a -> MinPQueue k a
MinPQueue.insert Node a b c
node.value Node a b c
node Frontier a b c
frontier.queue,
$sel:indexSetA:Frontier :: IntSet
indexSetA = Int -> IntSet -> IntSet
IntSet.insert Int
ia Frontier a b c
frontier.indexSetA,
$sel:indexSetB:Frontier :: IntSet
indexSetB = Int -> IntSet -> IntSet
IntSet.insert Int
ib Frontier a b c
frontier.indexSetB
}
mkNode :: (a -> b -> c) -> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode :: forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
ia, Int
ib) NonEmpty a
as NonEmpty b
bs =
Node
{ $sel:position:Node :: (Int, Int)
position = (Int
ia, Int
ib),
$sel:value:Node :: c
value = a -> b -> c
f (NonEmpty a -> a
forall a. NonEmpty a -> a
NonEmpty.head NonEmpty a
as) (NonEmpty b -> b
forall a. NonEmpty a -> a
NonEmpty.head NonEmpty b
bs),
$sel:as:Node :: NonEmpty a
as = NonEmpty a
as,
$sel:bs:Node :: NonEmpty b
bs = NonEmpty b
bs
}