51 lines
1.3 KiB
Haskell
51 lines
1.3 KiB
Haskell
{-# LANGUAGE NamedFieldPuns #-}
|
|
module Priority (
|
|
Queue(..)
|
|
, empty
|
|
, set
|
|
, toList
|
|
) where
|
|
|
|
import Data.Map (Map)
|
|
import qualified Data.Map as Map (
|
|
adjust, alter, delete, empty, insert, lookup, lookupMax
|
|
)
|
|
|
|
data Queue a = Queue {
|
|
byElement :: Map a Int
|
|
, byRank :: Map Int [a]
|
|
}
|
|
|
|
empty :: Queue a
|
|
empty = Queue {byElement = Map.empty, byRank = Map.empty}
|
|
|
|
set :: Ord a => a -> Int -> Queue a -> Queue a
|
|
set x priority queue@(Queue {byElement, byRank}) =
|
|
case Map.lookup x byElement of
|
|
Nothing -> Queue {
|
|
byElement = Map.insert x priority byElement
|
|
, byRank = Map.alter (push x) priority byRank
|
|
}
|
|
Just formerPriority
|
|
| formerPriority >= priority -> queue
|
|
| otherwise -> queue {
|
|
byElement = Map.insert x priority byElement
|
|
, byRank = move x formerPriority priority byRank
|
|
}
|
|
|
|
push :: a -> Maybe [a] -> Maybe [a]
|
|
push a Nothing = Just [a]
|
|
push a (Just as) = Just (a:as)
|
|
|
|
move :: Eq a => a -> Int -> Int -> Map Int [a] -> Map Int [a]
|
|
move a from to =
|
|
Map.alter (push a) to . Map.adjust (filter (/= a)) from
|
|
|
|
toList :: Queue a -> [(Int, a)]
|
|
toList = extract . byRank
|
|
where
|
|
extract tmpMap =
|
|
case Map.lookupMax tmpMap of
|
|
Nothing -> []
|
|
Just (k, elems) -> ((,) k <$> elems) ++ extract (Map.delete k tmpMap)
|