src/task9/task9.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Data.List (inits, sortBy, foldl')
   4 import Data.Function (on)
   5
   6 import Data.Text (Text)
   7 import qualified Data.Text.IO as TIO
   8
   9 import Data.Void (Void)
  10
  11 import Text.Megaparsec
  12 import Text.Megaparsec.Char
  13 import qualified Text.Megaparsec.Char.Lexer as L
  14 import qualified Control.Applicative as CA
  15
  16 import qualified Data.HashMap.Strict as M
  17 import Data.HashMap.Strict ((!))
  18
  19
  20 type Weight = Int
  21 type Value = Int
  22 data Bag = Bag { bagWeight :: Weight, bagValue :: Value} deriving (Show, Eq)
  23 data TableEntry = TableEntry {value :: Value, backpointer :: (Int, Weight)} deriving (Show, Eq)
  24 type DPTable = M.HashMap (Int, Weight) TableEntry
  25
  26 main :: IO ()
  27 main = do
  28         bagsT <- TIO.readFile "data/09-bags.txt"
  29         let bags = successfulParse bagsT
  30         let limit = 5000
  31         print $ part1 limit bags
  32         let dpTable = buildTable limit bags
  33         print $ part2 dpTable limit bags
  34         print $ backtrace dpTable ((length bags), limit)
  35
  36 part1 :: Weight -> [Bag] -> Int
  37 part1 limit bags = length $ last $ takeWhile (willFit limit) $ inits sortedBags
  38     where sortedBags = sortBy (compare `on` bagWeight) bags
  39           willFit limit' bags' = (sum $ map bagWeight bags') < limit'
  40
  41 part2 :: DPTable -> Weight -> [Bag] -> Value
  42 part2 table limit bags = value $ table!(length bags, limit)
  43
  44
  45 buildTable :: Weight -> [Bag] -> DPTable
  46 buildTable limit bags = table
  47     where initialTable = M.fromList [((0, w), TableEntry {value = 0, backpointer = (0, w)}) | w <- [0..limit]]
  48           table = foldl' includeBag initialTable
  49                           [ (thisBag, bagNumber, allowedWeight)
  50                           | (thisBag, bagNumber) <- zip bags [1..]
  51                           , allowedWeight <- [0..limit]
  52                           ]
  53
  54
  55 includeBag :: DPTable -> (Bag, Int, Weight) -> DPTable
  56 includeBag table (bag, bagNumber, weight) =
  57     if bagWeight bag > weight
  58     then M.insert here withoutBag table
  59     else if value withBagEntry + bagValue bag > value withoutBagEntry
  60          then M.insert here withBag table
  61          else M.insert here withoutBag table
  62     where here = (bagNumber, weight)
  63           withoutBagKey = (bagNumber - 1, weight)
  64           withBagKey = (bagNumber - 1, weight - bagWeight bag)
  65           withoutBagEntry = table!withoutBagKey
  66           withBagEntry = table!withBagKey
  67           withoutBag = TableEntry {value = value withoutBagEntry, backpointer = withoutBagKey}
  68           withBag = TableEntry {value = value withBagEntry + bagValue bag, backpointer = withBagKey}
  69
  70
  71 backtrace :: DPTable -> (Int, Weight) -> [Bag]
  72 backtrace table key@(numberBags, weight)
  73     | numberBags == 0       = [Bag { bagValue = 0, bagWeight = weight }]
  74     | weight == nextWeight  = backtrace table nextKey
  75     | otherwise             = (backtrace table nextKey) ++ [includedBag]
  76     where nextKey = backpointer (table!key)
  77           (_, nextWeight) = nextKey
  78           includedBag = Bag { bagWeight = weight - nextWeight
  79                             , bagValue = (value (table!key)) - value (table!nextKey)}
  80
  81 -- Parse the input file
  82
  83 type Parser = Parsec Void Text
  84
  85 -- don't treat newlines as automatically-consumed whitespace
  86 sc :: Parser ()
  87 sc = L.space (skipSome (char ' ')) CA.empty CA.empty
  88
  89 lexeme  = L.lexeme sc
  90 integer = lexeme L.decimal
  91 -- symb = L.symbol sc
  92
  93 bagsFileP = bagP `sepBy` newline
  94
  95 bagP = bagify <$> integer <*> integer
  96     where bagify w v = Bag { bagWeight = w , bagValue = v}
  97
  98 successfulParse :: Text -> [Bag]
  99 successfulParse input =
 100         case parse bagsFileP "input" input of
 101                 Left  _error -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 102                 Right bags -> bags