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 bags
  32         print $ part1 limit bags
  33         print $ part2 limit bags
  34
  35 part1 :: Weight -> [Bag] -> Int
  36 part1 limit bags = length $ last $ takeWhile (willFit limit) $ inits sortedBags
  37     where sortedBags = sortBy (compare `on` bagWeight) bags
  38           willFit limit' bags' = (sum $ map bagWeight bags') < limit'
  39
  40 part2 :: Weight -> [Bag] -> Value
  41 part2 limit bags = value $ table!(length bags, limit)
  42     where initialTable = M.fromList [((0, w), TableEntry {value = 0, backpointer = (0, w)}) | w <- [0..limit]]
  43           table = foldl' includeBag initialTable
  44                           [ (thisBag, bagNumber, allowedWeight)
  45                           | (thisBag, bagNumber) <- zip bags [1..]
  46                           , allowedWeight <- [0..limit]
  47                           ]
  48
  49
  50 includeBag :: DPTable -> (Bag, Int, Weight) -> DPTable
  51 includeBag table (bag, bagNumber, weight) =
  52     if bagWeight bag > weight
  53     then M.insert here withoutBag table
  54     else if value withBagEntry + bagValue bag > value withoutBagEntry
  55          then M.insert here withBag table
  56          else M.insert here withoutBag table
  57     where here = (bagNumber, weight)
  58           withoutBagKey = (bagNumber - 1, weight)
  59           withBagKey = (bagNumber - 1, weight - bagWeight bag)
  60           withoutBagEntry = table!withoutBagKey
  61           withBagEntry = table!withBagKey
  62           withoutBag = TableEntry {value = value withoutBagEntry, backpointer = withoutBagKey}
  63           withBag = TableEntry {value = value withBagEntry + bagValue bag, backpointer = withBagKey}
  64
  65
  66 -- Parse the input file
  67
  68 type Parser = Parsec Void Text
  69
  70 -- don't treat newlines as automatically-consumed whitespace
  71 sc :: Parser ()
  72 sc = L.space (skipSome (char ' ')) CA.empty CA.empty
  73
  74 lexeme  = L.lexeme sc
  75 integer = lexeme L.decimal
  76 -- symb = L.symbol sc
  77
  78 bagsFileP = bagP `sepBy` newline
  79
  80 bagP = bagify <$> integer <*> integer
  81     where bagify w v = Bag { bagWeight = w , bagValue = v}
  82
  83 successfulParse :: Text -> [Bag]
  84 successfulParse input =
  85         case parse bagsFileP "input" input of
  86                 Left  _error -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  87                 Right bags -> bags