diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,3 @@ +## Assignment 0.0.1.0 + +* Initial release. diff --git a/Data/Algorithm/Assignment.hs b/Data/Algorithm/Assignment.hs new file mode 100644 --- /dev/null +++ b/Data/Algorithm/Assignment.hs @@ -0,0 +1,327 @@ +-- | +-- Module : Data.Algorithm.Assignment +-- Copyright : © 2024–present Mark Karpov +-- License : BSD 3 clause +-- +-- Maintainer : Mark Karpov +-- Stability : experimental +-- Portability : portable +-- +-- A solution to the assignment problem. +module Data.Algorithm.Assignment + ( assign, + ) +where + +import Control.Monad (forM_, void, when) +import Control.Monad.Fix (fix) +import Control.Monad.ST (ST, runST) +import Data.Array (Array) +import Data.Array.Base qualified as A +import Data.Array.ST (STUArray) +import Data.Array.ST qualified as ST +import Data.STRef (modifySTRef', newSTRef, readSTRef, writeSTRef) + +type CostMatrix s = STUArray s (Int, Int) Int + +type MarkMatrix s = STUArray s (Int, Int) Char + +type CoverageVector s = STUArray s Int Bool + +-- | \(\mathcal{O}(n^4)\). Assign elements from two collections to each +-- other so that the total cost is minimal. The cost of each combination is +-- given the by the first argument and it can be negative. If any of the +-- collections is empty the result is the empty list. The sizes of the +-- collections need not to match. Finally, there is no guarantees on the +-- order of elements in the returned list of pairs. +-- +-- See: +assign :: + -- | How to calculate the cost + (a -> b -> Int) -> + -- | The first collection + [a] -> + -- | The second collection + [b] -> + -- | The resulting optimal assignment (no guarantees about order) + [(a, b)] +assign _ [] _ = [] +assign _ _ [] = [] +assign cost as bs = runST $ do + let length_a = length as + length_b = length bs + aMinBound = 0 + aMaxBound = length_a - 1 + bMinBound = 0 + bMaxBound = length_b - 1 + abMaxBound = max aMaxBound bMaxBound + asArray = A.listArray (aMinBound, aMaxBound) as + bsArray = A.listArray (bMinBound, bMaxBound) bs + matrixBounds = ((aMinBound, bMinBound), (abMaxBound, abMaxBound)) + c <- ST.newArray matrixBounds 0 + m <- ST.newArray matrixBounds noMark + aCoverage <- ST.newArray (aMinBound, abMaxBound) False + bCoverage <- ST.newArray (bMinBound, abMaxBound) False + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> + ST.writeArray c (i, j) (cost (asArray A.! i) (bsArray A.! j)) + if aMaxBound - aMinBound >= bMaxBound - bMinBound + then normalizePerB c + else normalizePerA c + starZeros c m aCoverage bCoverage + fix $ \recurse0 -> do + done <- coverZeros m aCoverage bCoverage + if done + then recoverResults m asArray bsArray + else fix $ \recurse1 -> do + r <- primeUncoveredZero c m aCoverage bCoverage + case r of + Nothing -> do + adjustCosts c aCoverage bCoverage + recurse1 + Just z0 -> do + adjustMarks m z0 + clearCoverage aCoverage bCoverage + recurse0 +{-# INLINEABLE assign #-} + +normalizePerA :: CostMatrix s -> ST s () +normalizePerA c = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds c + countFromTo aMinBound aMaxBound $ \i -> do + minValueRef <- newSTRef maxBound + countFromTo bMinBound bMaxBound $ \j -> + ST.readArray c (i, j) >>= modifySTRef' minValueRef . min + minValue <- readSTRef minValueRef + when (minValue /= 0) $ do + countFromTo bMinBound bMaxBound $ \j -> do + ST.modifyArray' c (i, j) (subtract minValue) +{-# INLINE normalizePerA #-} + +normalizePerB :: CostMatrix s -> ST s () +normalizePerB c = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds c + countFromTo bMinBound bMaxBound $ \j -> do + minValueRef <- newSTRef maxBound + countFromTo aMinBound aMaxBound $ \i -> + ST.readArray c (i, j) >>= modifySTRef' minValueRef . min + minValue <- readSTRef minValueRef + when (minValue /= 0) $ do + countFromTo aMinBound aMaxBound $ \i -> do + ST.modifyArray' c (i, j) (subtract minValue) +{-# INLINE normalizePerB #-} + +starZeros :: + CostMatrix s -> + MarkMatrix s -> + CoverageVector s -> + CoverageVector s -> + ST s () +starZeros c m aCoverage bCoverage = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds c + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> do + x <- ST.readArray c (i, j) + when (x == 0) $ do + aCovered <- ST.readArray aCoverage i + bCovered <- ST.readArray bCoverage j + when (not aCovered && not bCovered) $ do + ST.writeArray m (i, j) starMark + ST.writeArray aCoverage i True + ST.writeArray bCoverage j True + clearCoverage aCoverage bCoverage +{-# INLINE starZeros #-} + +coverZeros :: + MarkMatrix s -> + CoverageVector s -> + CoverageVector s -> + ST s Bool +coverZeros m _aCoverage bCoverage = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds m + nRef <- newSTRef 0 + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> do + x <- ST.readArray m (i, j) + bCovered <- ST.readArray bCoverage j + when (x == starMark && not bCovered) $ do + ST.writeArray bCoverage j True + modifySTRef' nRef (+ 1) + n <- readSTRef nRef + return (n > aMaxBound) +{-# INLINE coverZeros #-} + +recoverResults :: + MarkMatrix s -> + Array Int a -> + Array Int b -> + ST s [(a, b)] +recoverResults m as bs = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds m + resultRef <- newSTRef [] + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> do + x <- ST.readArray m (i, j) + when (x == starMark) $ do + case (,) <$> (as A.!? i) <*> (bs A.!? j) of + Nothing -> return () + Just (a, b) -> modifySTRef' resultRef ((a, b) :) + readSTRef resultRef +{-# INLINE recoverResults #-} + +primeUncoveredZero :: + CostMatrix s -> + MarkMatrix s -> + CoverageVector s -> + CoverageVector s -> + ST s (Maybe (Int, Int)) +primeUncoveredZero c m aCoverage bCoverage = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds m + primedRef <- newSTRef Nothing + void . countFromTo' aMinBound aMaxBound $ \i -> + countFromTo' bMinBound bMaxBound $ \j -> do + x <- ST.readArray c (i, j) + if x == 0 + then do + aCovered <- ST.readArray aCoverage i + bCovered <- ST.readArray bCoverage j + if not aCovered && not bCovered + then False <$ writeSTRef primedRef (Just (i, j)) + else return True + else return True + r <- readSTRef primedRef + case r of + Nothing -> return Nothing + Just (i, j) -> do + ST.writeArray m (i, j) primeMark + mj' <- findInA m starMark i + case mj' of + Nothing -> return (Just (i, j)) + Just j' -> do + ST.writeArray aCoverage i True + ST.writeArray bCoverage j' False + primeUncoveredZero c m aCoverage bCoverage +{-# INLINEABLE primeUncoveredZero #-} + +adjustMarks :: MarkMatrix s -> (Int, Int) -> ST s () +adjustMarks m z0 = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds m + let go (_, j) acc = do + mi' <- findInB m starMark j + case mi' of + Nothing -> return acc + Just i' -> do + mj' <- findInA m primeMark i' + case mj' of + Nothing -> error "Data.Algorithm.Assignment.adjustMarks" + Just j' -> go (i', j') ((i', j) : (i', j') : acc) + path <- go z0 [z0] + forM_ path $ \(i, j) -> do + let adjust x = + if x == starMark + then noMark + else starMark + ST.modifyArray' m (i, j) adjust + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> do + let resetPrime x = + if x == primeMark + then noMark + else x + ST.modifyArray' m (i, j) resetPrime +{-# INLINE adjustMarks #-} + +adjustCosts :: + CostMatrix s -> + CoverageVector s -> + CoverageVector s -> + ST s () +adjustCosts c aCoverage bCoverage = do + ((aMinBound, bMinBound), (aMaxBound, bMaxBound)) <- ST.getBounds c + minUncoveredValueRef <- newSTRef maxBound + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> do + aCovered <- ST.readArray aCoverage i + bCovered <- ST.readArray bCoverage j + when (not aCovered && not bCovered) $ do + ST.readArray c (i, j) >>= modifySTRef' minUncoveredValueRef . min + minUncoveredValue <- readSTRef minUncoveredValueRef + countFromTo aMinBound aMaxBound $ \i -> + countFromTo bMinBound bMaxBound $ \j -> do + aCovered <- ST.readArray aCoverage i + bCovered <- ST.readArray bCoverage j + if not aCovered && not bCovered + then ST.modifyArray c (i, j) (subtract minUncoveredValue) + else + when (aCovered && bCovered) $ + ST.modifyArray c (i, j) (+ minUncoveredValue) +{-# INLINE adjustCosts #-} + +clearCoverage :: + CoverageVector s -> + CoverageVector s -> + ST s () +clearCoverage aCoverage bCoverage = do + let clearOne v = do + (from, to) <- ST.getBounds v + countFromTo from to $ \i -> + ST.writeArray v i False + clearOne aCoverage + clearOne bCoverage +{-# INLINE clearCoverage #-} + +findInA :: + MarkMatrix s -> + Char -> + Int -> + ST s (Maybe Int) +findInA m mark i = do + ((_aMinBound, bMinBound), (_aMaxBound, bMaxBound)) <- ST.getBounds m + starredRef <- newSTRef Nothing + void . countFromTo' bMinBound bMaxBound $ \j -> do + x <- ST.readArray m (i, j) + if x == mark + then False <$ writeSTRef starredRef (Just j) + else return True + readSTRef starredRef +{-# INLINE findInA #-} + +findInB :: + MarkMatrix s -> + Char -> + Int -> + ST s (Maybe Int) +findInB m mark j = do + ((aMinBound, _bMinBound), (aMaxBound, _bMaxBound)) <- ST.getBounds m + starredRef <- newSTRef Nothing + void . countFromTo' aMinBound aMaxBound $ \i -> do + x <- ST.readArray m (i, j) + if x == mark + then False <$ writeSTRef starredRef (Just i) + else return True + readSTRef starredRef +{-# INLINE findInB #-} + +countFromTo :: Int -> Int -> (Int -> ST s ()) -> ST s () +countFromTo start end action = go start + where + go !n = when (n <= end) $ do + action n + go (n + 1) +{-# INLINE countFromTo #-} + +countFromTo' :: Int -> Int -> (Int -> ST s Bool) -> ST s Bool +countFromTo' start end action = go start + where + go !n = + if n <= end + then do + r <- action n + if r then go (n + 1) else return False + else return True +{-# INLINE countFromTo' #-} + +noMark, starMark, primeMark :: Char +noMark = 'n' +starMark = 's' +primeMark = 'p' diff --git a/LICENSE.md b/LICENSE.md new file mode 100644 --- /dev/null +++ b/LICENSE.md @@ -0,0 +1,28 @@ +Copyright © 2024–present Mark Karpov + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +* Neither the name Mark Karpov nor the names of contributors may be used to + endorse or promote products derived from this software without specific + prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS “AS IS” AND ANY EXPRESS +OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES +OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN +NO EVENT SHALL THE COPYRIGHT HOLDERS BE LIABLE FOR ANY DIRECT, INDIRECT, +INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, +EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,27 @@ +# Assignment + +[![License BSD3](https://img.shields.io/badge/license-BSD3-brightgreen.svg)](http://opensource.org/licenses/BSD-3-Clause) +[![Hackage](https://img.shields.io/hackage/v/assignment.svg?style=flat)](https://hackage.haskell.org/package/assignment) +[![Stackage Nightly](http://stackage.org/package/assignment/badge/nightly)](http://stackage.org/nightly/package/assignment) +[![Stackage LTS](http://stackage.org/package/assignment/badge/lts)](http://stackage.org/lts/package/assignment) +[![CI](https://github.com/mrkkrp/assignment/actions/workflows/ci.yaml/badge.svg)](https://github.com/mrkkrp/assignment/actions/workflows/ci.yaml) + +This package implements a solution to the [assignment problem][problem]. It +follows the [matrix interpretation][matrix-interpretation] of the so-called +Hungarian algorithm and works in *O(n^4)*. + +[problem]: https://en.wikipedia.org/wiki/Assignment_problem +[matrix-interpretation]: https://en.wikipedia.org/wiki/Hungarian_algorithm#Matrix_interpretation + +## Contribution + +Issues, bugs, and questions may be reported in [the GitHub issue tracker for +this project](https://github.com/mrkkrp/assignment/issues). + +Pull requests are also welcome. + +## License + +Copyright © 2024–present Mark Karpov + +Distributed under BSD 3 clause license. diff --git a/assignment.cabal b/assignment.cabal new file mode 100644 --- /dev/null +++ b/assignment.cabal @@ -0,0 +1,96 @@ +cabal-version: 2.4 +name: assignment +version: 0.0.1.0 +license: BSD-3-Clause +license-file: LICENSE.md +maintainer: Mark Karpov +author: Mark Karpov +tested-with: ghc ==9.8.2 ghc ==9.10.1 +homepage: https://github.com/mrkkrp/assignment +bug-reports: https://github.com/mrkkrp/assignment/issues +synopsis: A solution to the assignment problem +description: A solution to the assignment problem. +category: Algorithms +build-type: Simple +extra-doc-files: + CHANGELOG.md + README.md + +source-repository head + type: git + location: https://github.com/mrkkrp/assignment.git + +flag dev + description: Turn on development settings. + default: False + manual: True + +library + exposed-modules: Data.Algorithm.Assignment + default-language: GHC2021 + build-depends: + base >=4.15 && <5, + array >=0.5.6 && <0.6 + + if flag(dev) + ghc-options: + -Wall -Werror -Wredundant-constraints -Wpartial-fields + -Wunused-packages -Wno-unused-imports + + else + ghc-options: -O2 -Wall + +test-suite tests + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: tests + default-language: GHC2021 + build-depends: + QuickCheck >=2.14, + assignment, + base >=4.15 && <5, + hspec >=2.0 && <3 + + if flag(dev) + ghc-options: + -Wall -Werror -Wredundant-constraints -Wpartial-fields + -Wunused-packages + + else + ghc-options: -O2 -Wall + +benchmark bench-speed + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench/speed + default-language: GHC2021 + build-depends: + assignment, + base >=4.15 && <5, + criterion >=0.6.2.1 && <1.7 + + if flag(dev) + ghc-options: + -Wall -Werror -Wredundant-constraints -Wpartial-fields + -Wunused-packages + + else + ghc-options: -O2 -Wall + +benchmark bench-memory + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench/memory + default-language: GHC2021 + build-depends: + assignment, + base >=4.15 && <5, + weigh >=0.0.4 + + if flag(dev) + ghc-options: + -Wall -Werror -Wredundant-constraints -Wpartial-fields + -Wunused-packages + + else + ghc-options: -O2 -Wall diff --git a/bench/memory/Main.hs b/bench/memory/Main.hs new file mode 100644 --- /dev/null +++ b/bench/memory/Main.hs @@ -0,0 +1,25 @@ +module Main (main) where + +import Data.Algorithm.Assignment (assign) +import Weigh + +main :: IO () +main = mainWith $ do + setColumns [Case, Allocated, GCs, Max] + benchCase 5 5 + benchCase 5 10 + benchCase 10 5 + benchCase 10 10 + benchCase 10 20 + benchCase 20 10 + benchCase 20 20 + benchCase 40 40 + +benchCase :: Int -> Int -> Weigh () +benchCase a b = + func name (uncurry (assign cost)) (as, bs) + where + name = "assign " ++ show a ++ " " ++ show b + as = [1 .. a] + bs = [1 .. b] + cost x y = x `div` y - x `rem` y + (y * 3) `rem` 4 diff --git a/bench/speed/Main.hs b/bench/speed/Main.hs new file mode 100644 --- /dev/null +++ b/bench/speed/Main.hs @@ -0,0 +1,26 @@ +module Main (main) where + +import Criterion.Main +import Data.Algorithm.Assignment (assign) + +main :: IO () +main = + defaultMain + [ benchCase 5 5, + benchCase 5 10, + benchCase 10 5, + benchCase 10 10, + benchCase 10 20, + benchCase 20 10, + benchCase 20 20, + benchCase 40 40 + ] + +benchCase :: Int -> Int -> Benchmark +benchCase a b = + env (return (as, bs)) (bench name . nf (uncurry (assign cost))) + where + name = "assign " ++ show a ++ " " ++ show b + as = [1 .. a] + bs = [1 .. b] + cost x y = x `div` y - x `rem` y + (y * 3) `rem` 4 diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -0,0 +1,59 @@ +module Main (main) where + +import Data.Algorithm.Assignment (assign) +import Data.List qualified +import Data.Tuple (swap) +import Test.Hspec +import Test.QuickCheck + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + describe "assign" $ do + it "returns the empty list if the first collection is empty" $ + assign (\_ _ -> 0) ([] :: [Int]) ([1, 2, 3] :: [Int]) `shouldBe` [] + it "returns the empty list if the second collection is empty" $ + assign (\_ _ -> 0) ([1, 2, 3] :: [Int]) ([] :: [Int]) `shouldBe` [] + it "solves the example from Wikipedia" $ do + let cost "Alice" "Clean bathroom" = 8 + cost "Alice" "Sweep floors" = 4 + cost "Alice" "Wash windows" = 7 + cost "Bob" "Clean bathroom" = 5 + cost "Bob" "Sweep floors" = 2 + cost "Bob" "Wash windows" = 3 + cost "Dora" "Clean bathroom" = 9 + cost "Dora" "Sweep floors" = 4 + cost "Dora" "Wash windows" = 8 + cost _ _ = 0 + Data.List.sortOn + fst + ( assign + cost + ["Alice", "Bob", "Dora"] + ["Clean bathroom", "Sweep floors", "Wash windows"] + ) + `shouldBe` [ ("Alice", "Clean bathroom"), + ("Bob", "Wash windows"), + ("Dora", "Sweep floors") + ] + it "always finds the best assignment" $ + property $ \(Fun _ cost) (as :: [Int]) (bs :: [Int]) -> do + let totalCost = sum . fmap cost + forAll (shuffle bs) $ \bs' -> + totalCost (assign (curry cost) as bs) <= totalCost (zip as bs') + it "swapping of the collections results in pairings with the same total cost" $ + property $ \(Fun _ cost) (as :: [Int]) (bs :: [Int]) -> do + let cost0 = curry cost + cost1 x y = cost0 y x + totalCost = sum . fmap cost + totalCost (assign cost0 as bs) == totalCost (swap <$> assign cost1 bs as) + it "shuffling of collections results in pairings with the same total cost" $ + property $ \(Fun _ cost) (xs :: [(Int, Int)]) -> do + let totalCost = sum . fmap cost + cost' = curry cost + as = Data.List.nub (fst <$> xs) + bs = Data.List.nub (snd <$> xs) + forAll (shuffle as) $ \as' -> + totalCost (assign cost' as bs) == totalCost (assign cost' as' bs)