module Data.Multimap (
Multimap,
null,
size,
numKeys,
numValues,
member,
notMember,
lookup,
(!),
empty,
insert,
delete,
map,
mapKeys,
mapWithKey,
foldr,
foldl,
foldrWithKey,
foldlWithKey,
elems,
keys,
keysSet,
assocs,
toMap,
toMapOfSets,
toList,
fromList,
findMin,
findMax,
findMinWithValues,
findMaxWithValues
) where
import Prelude hiding (lookup, map, null, foldr, foldl)
import qualified Prelude as P
import qualified Data.Set as Set
import Data.Set (Set)
import qualified Data.Map as Map
import Data.Map (Map)
import Data.Word
newtype Multimap k v = Multimap (Word32, Word32, Map k [v])
null :: Multimap k a -> Bool
null (Multimap (_, _, m)) = Map.null m
size :: Multimap k a -> Int
size (Multimap (_, size, _)) = fromIntegral size
numKeys :: Multimap k a -> Word32
numKeys (Multimap (nk, _, _)) = nk
numValues :: Multimap k a -> Word32
numValues (Multimap (_, nv, _)) = nv
notMember, member :: Ord k => Multimap k a -> k -> Bool
member (Multimap (_, _, map)) key = Map.member key map
notMember key = not . member key
(!) :: Ord k => k -> Multimap k a -> [a]
(!) = lookup
lookup :: Ord k => k -> Multimap k a -> [a]
lookup key (Multimap (_, _, map)) = maybe [] id (Map.lookup key map)
empty :: Multimap k a
empty = Multimap (0, 0, Map.empty)
insert :: Ord k => k -> a -> Multimap k a -> Multimap k a
insert k v (Multimap (nk, nv, map))
| Map.member k map = Multimap (nk, succ nv, Map.insert k (v : map Map.! k) map)
| otherwise = Multimap (succ nk, succ nv, Map.insert k [v] map)
delete :: Ord k => k -> a -> Multimap k a -> Multimap k a
delete k v m@(Multimap (nk, nv, map)) = case Map.lookup k map of
Just v -> Multimap (pred nk, nv fromIntegral (length v), Map.delete k map)
_ -> m
map :: (a -> b) -> Multimap k a -> Multimap k b
map f (Multimap (nk, nv, map)) = Multimap (nk, nv, Map.map (P.map f) map)
mapKeys :: Ord k2 => (k1 -> k2) -> Multimap k1 a -> Multimap k2 a
mapKeys f (Multimap (nk, nv, map)) = Multimap (nk, nv, Map.mapKeys f map)
mapWithKey :: (k -> a -> b) -> Multimap k a -> Multimap k b
mapWithKey f (Multimap (nk, nv, map))
= Multimap (nk, nv, Map.mapWithKey (\k -> P.map (f k)) map)
foldr :: (a -> b -> b) -> b -> Multimap k a -> b
foldr f e = P.foldr f e . concat . elems
foldl :: (a -> b -> a) -> a -> Multimap k b -> a
foldl f e = P.foldl f e . concat . elems
foldrWithKey :: (k -> a -> b -> b) -> b -> Multimap k a -> b
foldrWithKey f e = P.foldr (uncurry f) e . toList
foldlWithKey :: (a -> k -> b -> a) -> a -> Multimap k b -> a
foldlWithKey f e = P.foldl (\a (k,v) -> f a k v) e . toList
elems :: Multimap k a -> [[a]]
elems (Multimap (_, _, map)) = Map.elems map
keys :: Multimap k a -> [k]
keys (Multimap (_, _, map)) = Map.keys map
keysSet :: Multimap k a -> Set k
keysSet (Multimap (_, _, map)) = Map.keysSet map
assocs :: Multimap k a -> [(k, [a])]
assocs (Multimap (_, _, map)) = Map.assocs map
toMap :: Multimap k a -> Map k [a]
toMap (Multimap (_, _, theUnderlyingMap)) = theUnderlyingMap
toMapOfSets :: Ord a => Multimap k a -> Map k (Set a)
toMapOfSets (Multimap (_, _, map)) = Map.map Set.fromList map
toList :: Multimap k a -> [(k, a)]
toList (Multimap (_, _, map))
= concat $ Map.elems $ Map.mapWithKey (\k -> zip (repeat k)) map
fromList :: Ord k => [(k, a)] -> Multimap k a
fromList = P.foldr (uncurry insert) empty
findMin :: Multimap k a -> Maybe k
findMin (Multimap (_, _, map))
| Map.null map = Nothing
| otherwise = Just $ fst $ Map.findMin map
findMax :: Multimap k a -> Maybe k
findMax (Multimap (_, _, map))
| Map.null map = Nothing
| otherwise = Just $ fst $ Map.findMax map
findMinWithValues :: Multimap k a -> Maybe (k, [a])
findMinWithValues (Multimap (_, _, map))
| Map.null map = Nothing
| otherwise = Just $ Map.findMin map
findMaxWithValues :: Multimap k a -> Maybe (k, [a])
findMaxWithValues (Multimap (_, _, map))
| Map.null map = Nothing
| otherwise = Just $ Map.findMax map