src/advent15/advent15.hs

   1 import Data.List
   2 import Data.Tuple (swap)
   3 import Data.Maybe
   4 -- import Data.Foldable (forM_)
   5 import qualified Data.Map.Strict as M
   6 import Data.Map.Strict ((!))
   7 import qualified Data.Set as S
   8
   9 type Coord = (Int, Int) -- row, column
  10 type HitPoints = Int
  11 data Species = Elf | Goblin deriving (Show, Eq)
  12 data Agent = Agent Species HitPoints deriving (Show, Eq)
  13 type Agents = M.Map Coord Agent
  14 type Layout = S.Set Coord
  15 type Distances = M.Map Coord Int
  16
  17 newGoblin = Agent Goblin 200
  18 newElf = Agent Elf 200
  19
  20 isSpecies s (Agent s' _) = s == s'
  21
  22 isElf = isSpecies Elf
  23 isGoblin = isSpecies Goblin
  24
  25 otherSpecies Elf = Goblin
  26 otherSpecies Goblin = Elf
  27
  28 elfs :: Agents -> Agents
  29 elfs agents = M.filter isElf agents
  30
  31 goblins :: Agents -> Agents
  32 goblins agents = M.filter isGoblin agents
  33
  34 -- test1 = "#######\n#E..G.#\n#...#.#\n#.G.#G#\n#######\n"
  35 -- test2 = "#########\n#G..G..G#\n#.......#\n#.......#\n#G..E..G#\n#.......#\n#.......#\n#G..G..G#\n#########\n"
  36 -- test3 = "#######\n#.G...#\n#...EG#\n#.#.#G#\n#..G#E#\n#.....#\n#######\n"
  37 -- test4 = "#######\n#G..#E#\n#E#E.E#\n#G.##.#\n#...#E#\n#...E.#\n#######\n"
  38 -- test5 = "#######\n#E..EG#\n#.#G.E#\n#E.##E#\n#G..#.#\n#..E#.#\n#######\n"
  39 -- test6 = "#######\n#E.G#.#\n#.#G..#\n#G.#.G#\n#G..#.#\n#...E.#\n#######\n"
  40 -- test7 = "#######\n#.E...#\n#.#..G#\n#.###.#\n#E#G#G#\n#...#G#\n#######\n"
  41 -- test8 = "#########\n#G......#\n#.E.#...#\n#..##..G#\n#...##..#\n#...#...#\n#.G...G.#\n#.....G.#\n#########\n"
  42
  43 main :: IO ()
  44 main = do
  45     text <- readFile "data/advent15.txt"
  46     let (layout, agents) = parse text
  47     print $ part1 layout agents
  48     print $ part2 layout agents
  49  --   print layout
  50  --   print agents
  51     -- putStrLn $ showWorld layout agents
  52     -- let game = runGame layout agents 15
  53     -- print $ length game
  54     -- putStrLn $ showWorld layout $ snd $ last game
  55     -- print $ last game
  56     -- print $ length game - 1
  57     -- print $ scoreGame game
  58     -- print $ wonWithoutLoss game
  59     -- let game2 = runGame layout agents 14
  60     -- print $ length game2
  61     -- putStrLn $ showWorld layout $ snd $ last game2
  62     -- print $ last game2
  63     -- print $ length game2 - 1
  64     -- print $ scoreGame game2
  65     -- print $ wonWithoutLoss game2
  66     -- print $ bestMove (1, 1) layout agents
  67     -- let (h, e) = M.findMin $ elfs agents
  68     -- print (h, e)
  69     -- forM_ (M.assocs agents) $ \(ah, a) -> do
  70     --     print (ah, a)
  71         -- print $ pathsToEnemies ah layout agents
  72         -- print $ bestMove ah layout agents
  73     -- let a1 = doRound layout agents
  74     -- let a0 = doRound layout agents
  75     -- putStrLn $ "1\n" ++ showWorld layout a0
  76     -- print a0
  77     -- let n = 35
  78     -- let a1 = doNRounds n layout agents
  79     -- putStrLn $ show n ++ "\n" ++ showWorld layout a1
  80     -- print a1
  81     -- let a2 = doRound layout a1
  82     -- putStrLn $ show (n+1) ++ "\n" ++ showWorld layout a2
  83     -- print a2
  84     -- let a3 = doRound layout a2
  85     -- putStrLn $ show (n+2) ++ "\n" ++ showWorld layout a3
  86     -- print a3
  87     -- let a4 = doRound layout a3
  88     -- putStrLn $ show (n+3) ++ "\n" ++ showWorld layout a4
  89     -- print a4
  90     -- let a5 = doRound layout a4
  91     -- putStrLn $ show (n+4) ++ "\n" ++ showWorld layout a5
  92     -- print a5
  93
  94     -- let a6 = doRound layout a5
  95     -- putStrLn $ "27\n" ++ showWorld layout a6
  96     -- print a6
  97     -- let a7 = doRound layout a6
  98     -- putStrLn $ "28\n" ++ showWorld layout a7
  99     -- print a7
 100     -- let a8 = doNRounds 19 layout a7
 101     -- putStrLn $ "47\n" ++ showWorld layout a8
 102     -- print a8
 103     -- let a9 = doRound layout a8
 104     -- putStrLn $ "30\n" ++ showWorld layout a9
 105     -- print a9
 106     -- let aa = doRound layout a9
 107     -- putStrLn $ "31\n" ++ showWorld layout aa
 108     -- print aa
 109     -- let ab = doRound layout aa
 110     -- putStrLn $ "32\n" ++ showWorld layout ab
 111     -- print ab
 112     -- let ac = doRound layout ab
 113     -- putStrLn $ "33\n" ++ showWorld layout ac
 114     -- print ac
 115     -- let ad = doRound layout ac
 116     -- putStrLn $ "34\n" ++ showWorld layout ad
 117     -- print ad
 118     -- let ae = doRound layout ad
 119     -- putStrLn $ "35\n" ++ showWorld layout ae
 120     -- print ae
 121
 122
 123 part1 layout agents = scoreGame $ runGame layout agents 3
 124
 125 part2 layout agents = runPart2 layout agents 4
 126
 127 runPart2 layout agents elfPower =
 128     if wonWithoutLoss game
 129     then (scoreGame game, elfPower)
 130     else runPart2 layout agents (elfPower + 1)
 131     where game = runGame layout agents elfPower
 132
 133
 134 -- showWorld layout agents = unlines rows
 135 --     where rows = map (showRow layout agents) [minRow..maxRow]
 136 --           minRow = fst $ S.findMin layout'
 137 --           maxRow = fst $ S.findMax layout'
 138 --           layout' = S.union layout $ S.fromList $ M.keys agents
 139
 140 -- showRow layout agents row = map (\col -> showCell (row, col) layout agents) [minCol..maxCol]
 141 --     where minCol = minimum $ map snd $ S.toList layout'
 142 --           maxCol = maximum $ map snd $ S.toList layout'
 143 --           layout' = S.union layout $ S.fromList $ M.keys agents
 144
 145 -- showCell c layout agents =
 146 --     if c `M.member` agents
 147 --     then if isElf a then 'E' else 'G'
 148 --     else if c `S.member` layout then '.' else '\x2593'
 149 --     where a = agents!c
 150
 151
 152 -- Parsing
 153 parse :: String -> (Layout, Agents)
 154 parse text = foldl' parseRow (S.empty, M.empty) $ zip [0..] $ lines text
 155
 156 parseRow :: (Layout, Agents) -> (Int, String) -> (Layout, Agents)
 157 parseRow (layout, agents) (r, row) = foldl' parseCellWithY (layout, agents) $ zip [0..] row
 158     where parseCellWithY = parseCell r
 159
 160 parseCell :: Int -> (Layout, Agents) -> (Int, Char) -> (Layout, Agents)
 161 parseCell r (layout, agents) (c, cell) =
 162     let here = (r, c)
 163     in case cell of
 164             'G' -> (S.insert here layout, M.insert here newGoblin agents)
 165             'E' -> (S.insert here layout, M.insert here newElf agents)
 166             '.' -> (S.insert here layout, agents)
 167             _   -> (layout, agents)
 168
 169 -- Locations
 170
 171 adjacent :: Coord -> Layout -> Layout
 172 adjacent (r, c) layout = S.intersection layout
 173                          $ S.fromList [(r+1, c), (r-1, c), (r, c+1), (r, c-1)]
 174
 175 free :: Coord -> Layout -> Agents -> Bool
 176 free here layout agents = (here `S.member` layout) && (here `M.notMember` agents)
 177
 178 adjacentFree :: Coord -> Layout -> Agents -> Layout
 179 adjacentFree here layout agents = S.filter (\c -> free c layout agents) (adjacent here layout)
 180
 181 -- Move selection
 182
 183 distancesFrom :: Coord -> Layout -> Agents -> Distances
 184 distancesFrom here layout agents = distanceFlood layout agents [here] (M.singleton here 0)
 185
 186 distanceFlood :: Layout -> Agents -> [Coord] -> Distances -> Distances
 187 distanceFlood layout agents boundary distances
 188     | null boundary = distances
 189     | otherwise = distanceFlood layout agents newBoundary newDistances
 190     where current = head boundary
 191           currentCost = distances!current
 192           neighbours = filter (\c -> c `M.notMember` distances) $ S.toList $ adjacentFree current layout agents
 193           newDistances = foldl' (\m l -> M.insert l (currentCost + 1) m) distances neighbours
 194           newBoundary = nub $ (tail boundary) ++ neighbours
 195
 196
 197 shortestDistanceStepTo :: Coord -> Coord -> Layout -> Agents -> Maybe (Int, Coord)
 198 shortestDistanceStepTo here there layout agents =
 199     if M.null distanceSteps
 200     then Nothing
 201     else Just $ head $ sort $ map swap $ M.toList distanceSteps
 202     where distances = distancesFrom there layout agents
 203           steps = adjacentFree here layout agents
 204           distanceSteps = M.filterWithKey (\k _ -> k `S.member` steps) distances -- S.map (\c -> (distances!c, c)) steps
 205
 206 stepsTowardsEnemies :: Coord -> Layout -> Agents -> [(Int, Coord)]
 207 stepsTowardsEnemies here layout agents =
 208       catMaybes
 209     $ map (\e -> shortestDistanceStepTo here e layout agents)
 210     $ S.toList $ enemyLocations here agents
 211
 212 enemyLocations :: Coord -> Agents -> Layout
 213 enemyLocations here agents = S.fromList $ M.keys $ M.filter (isSpecies enemySpecies) agents
 214     where Agent thisSpecies _ = agents!here
 215           enemySpecies = otherSpecies thisSpecies
 216
 217 bestMove :: Coord -> Layout -> Agents -> Coord
 218 bestMove here layout agents =
 219     if null steps
 220     then here
 221     else snd $ head $ sort $ steps
 222     where steps = stepsTowardsEnemies here layout agents
 223
 224 makeMove :: Coord -> Coord -> Agents -> Agents
 225 makeMove here there agents = M.insert there agent $ M.delete here agents
 226     where agent = agents!here
 227
 228 -- Attacking
 229
 230 bestTarget :: Coord -> Layout -> Agents -> Coord
 231 bestTarget here layout agents = keyOfMinHP $ M.filterWithKey (\c _ -> c `S.member` enemies) agents
 232     where enemies = touchingEnemies here layout agents
 233
 234 attack :: Coord -> Agents -> Int -> Agents
 235 attack target agents elfPower = if hp > power
 236                        then M.insert target (Agent species (hp - power)) agents
 237                        else M.delete target agents
 238     where Agent species hp = agents!target
 239           power = if species == Goblin then elfPower else 3
 240
 241 keyOfMinHP :: Agents -> Coord -- Ord b => M.Map a b -> a
 242 keyOfMinHP m = fst $ M.foldrWithKey mergeKV (M.findMin m) m
 243     where mergeKV k (Agent s v) (bestK, (Agent sb bestV)) =
 244             if v < bestV then (k, (Agent s v)) else (bestK, (Agent sb bestV))
 245
 246 makeAttack :: Coord -> Layout -> Agents -> Int -> Agents
 247 makeAttack here layout agents elfPower = attack target agents elfPower
 248     where target = bestTarget here layout agents
 249
 250 -- Game loop
 251
 252 -- doNRounds :: Int -> Layout -> Agents -> Int -> Agents
 253 -- doNRounds n layout agents elfPower
 254 --     | n == 0 = agents
 255 --     | otherwise = doNRounds (n-1) layout (snd $ doRound layout agents elfPower) elfPower
 256
 257
 258 doRound :: Layout -> Agents -> Int -> (Bool, Agents)
 259 doRound layout agents elfPower = agents'
 260     where agents' = foldl' (\(_, a) h -> agentAction h layout a elfPower) (True, agents) $ M.keys agents
 261
 262 touchingEnemies :: Coord -> Layout -> Agents -> Layout
 263 touchingEnemies here layout agents = S.intersection neighbourhood enemies
 264     where neighbourhood = adjacent here layout
 265           enemies = enemyLocations here agents
 266
 267 agentAction :: Coord -> Layout -> Agents -> Int -> (Bool, Agents)
 268 agentAction here layout agents elfPower
 269     | (M.null (elfs agents)) || (M.null (goblins agents)) = (False, agents)
 270     | here `M.notMember` agents = (True, agents)
 271     | S.null $ enemyLocations here agents = (True, agents)
 272     | otherwise = (True, agents'')
 273     where targets = touchingEnemies here layout agents
 274           here' = if S.null targets
 275                   then bestMove here layout agents
 276                   else here
 277           agents' = makeMove here here' agents
 278           targets' = touchingEnemies here' layout agents'
 279           agents'' = if S.null targets'
 280                      then agents'
 281                      else makeAttack here' layout agents' elfPower
 282
 283 runGame :: Layout -> Agents -> Int -> [(Bool, Agents)]
 284 runGame layout agents elfPower = states ++ [doRound layout (snd $ last states) elfPower]
 285     where states = takeWhile (\(f, _) -> f == True) $ iterate (\(_, a) -> doRound layout a elfPower) (True, agents)
 286
 287 scoreGame :: [(Bool, Agents)] -> Int
 288 scoreGame states = (length states - 2) * hps
 289     where hps = sum $ map (\(Agent _ hp) -> hp) $ M.elems $ snd $ last states
 290
 291 wonWithoutLoss :: [(Bool, Agents)] -> Bool
 292 wonWithoutLoss game = startingElfs == finishingElfs
 293     where startingElfs = M.size $ elfs $ snd $ head game
 294           finishingElfs = M.size $ elfs $ snd $ last game