module Data.Algorithm.Munkres.Simple (
Problem,
problem,
Solution, associatedA, associatedB, associationList)
where
import qualified Data.Set as S (Set, size, toAscList)
import qualified Data.Bimap as B (Bimap, lookup, lookupR, toList, fromList)
import qualified Data.Array.IArray as IA (array)
import Data.Algorithm.Munkres (hungarianMethodDouble)
data Problem a b = Problem {
setA :: S.Set a
, setB :: S.Set b
, weightFunction :: a -> b -> Double
}
problem :: (Ord a, Ord b) => S.Set a -> S.Set b
-> (a -> b -> Double)
-> Problem a b
problem as bs df =
if (S.size as) == (S.size bs)
then Problem as bs df
else error "Sets are of different size"
data Solution a b = Solution {
solutionBimap :: B.Bimap a b
, solutionCost :: Double
}
solve :: (Ord a, Ord b) => Problem a b -> Solution a b
solve p = Solution solution cost
where
solution = B.fromList $
map (\(m,n) -> (snd $ as !! (m1), snd $ bs !! (n1))) intSolution
(intSolution, cost) = hungarianMethodDouble
$ IA.array ((1,1), (numA, numB)) assocs
numA = S.size . setA $ p
numB = S.size . setB $ p
as = zip [1..] (S.toAscList . setA $ p)
bs = zip [1..] (S.toAscList . setB $ p)
assocs = [((m,n), weightFunction p a b) | (m,a) <- as, (n, b) <- bs]
associatedA :: (Ord a, Ord b) => Solution a b -> b -> Maybe a
associatedA s b = B.lookupR b . solutionBimap $ s
associatedB :: (Ord a, Ord b) => Solution a b -> a -> Maybe b
associatedB s a = B.lookup a . solutionBimap $ s
associationList :: Solution a b -> [(a,b)]
associationList = B.toList . solutionBimap