-- |
-- Module      : MonusWeightedSearch.Examples.SubsetSum
-- Copyright   : (c) Donnacha Oisín Kidney 2021
-- Maintainer  : mail@doisinkidney.com
-- Stability   : experimental
-- Portability : non-portable
--
-- An implementation of shortest subset sum (the Inclusion-Exclusion method)
-- using the 'Heap' monad.

module MonusWeightedSearch.Examples.SubsetSum where

import Control.Monad.Heap
import Data.Monus.Dist
import Control.Monad.Writer
import Data.Maybe
import Control.Monad (filterM, guard)

-- | A weight for the inclusion or exclusion of an element.
--
-- This lets us weight the computation by number of elements included, and
-- therefore optimise for the fewest.
inclusion :: Monad m => HeapT Dist m Bool
inclusion :: forall (m :: Type -> Type). Monad m => HeapT Dist m Bool
inclusion = [(Bool, Dist)] -> HeapT Dist m Bool
forall (m :: Type -> Type) a w.
Applicative m =>
[(a, w)] -> HeapT w m a
fromList [(Bool
False,Dist
0),(Bool
True,Dist
1)]

-- | @'shortest' n xs@ returns the shortes subset of @xs@ which sums to @n@.
--
-- >>> shortest 5 [10,-4,3,11,6,12,1]
-- [-4,3,6]
shortest ::  Int -> [Int] -> [Int]
shortest :: Int -> [Int] -> [Int]
shortest Int
t [Int]
xs = (Dist, [Int]) -> [Int]
forall a b. (a, b) -> b
snd ((Dist, [Int]) -> [Int])
-> (Heap Dist [Int] -> (Dist, [Int])) -> Heap Dist [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (Dist, [Int]) -> (Dist, [Int])
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Dist, [Int]) -> (Dist, [Int]))
-> (Heap Dist [Int] -> Maybe (Dist, [Int]))
-> Heap Dist [Int]
-> (Dist, [Int])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Heap Dist [Int] -> Maybe (Dist, [Int])
forall w a. Monus w => Heap w a -> Maybe (w, a)
best (Heap Dist [Int] -> [Int]) -> Heap Dist [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ do
  [Int]
subset <- (Int -> HeapT Dist Identity Bool) -> [Int] -> Heap Dist [Int]
forall (m :: Type -> Type) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (HeapT Dist Identity Bool -> Int -> HeapT Dist Identity Bool
forall a b. a -> b -> a
const HeapT Dist Identity Bool
forall (m :: Type -> Type). Monad m => HeapT Dist m Bool
inclusion) [Int]
xs
  Bool -> HeapT Dist Identity ()
forall (f :: Type -> Type). Alternative f => Bool -> f ()
guard ([Int] -> Int
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
sum [Int]
subset Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t)
  pure [Int]
subset