aoc2020/14.hs

47 lines
1.6 KiB
Haskell
Raw Normal View History

2020-12-14 18:29:10 +00:00
import Data.Bits
import Data.Int
import Data.List
import Debug.Trace
import Data.List.Split
import Data.Maybe
import qualified Data.Map.Strict as Map
type State = (Int64, [Int], Map.Map Int64 Int64)
parseMask :: (Int64, [Int]) -> (Int, Char) -> (Int64, [Int])
parseMask (oMask, flucBits) (nth, ch) =
let oMask2 = if ch == '1' then oMask .|. (shift 1 nth) else oMask
flucBits2 = if ch == 'X' then nth : flucBits else flucBits in
(oMask2, flucBits2)
applyMask :: Bits a => a -> (Int, Int) -> a
applyMask n (i, c) = if c == 0 then clearBit n i else setBit n i
nextState :: State -> String -> State
nextState (oMask, flucBits, intMap) origLine =
if isPrefixOf "mask" origLine then
let line = fromJust $ stripPrefix "mask = " origLine in
let len = length line - 1 in
let enumLine = zip [len, len-1..0] line in
let (oMask, flucBits) = foldl parseMask (0, []) $ enumLine in
(oMask, flucBits, intMap)
else
let line = fromJust $ stripPrefix "mem[" origLine in
let [addr, value] = map read $ splitOn "] = " line in
let flucIter = sequence $ take (length flucBits) $ repeat [0, 1] in
let enumBits lines = zip flucBits lines in
let f = foldl applyMask (addr .|. oMask) . enumBits in
let addrs = map f flucIter in
let intMap2 = foldl (\acc el -> Map.insert el value acc) intMap addrs in
(oMask, flucBits, intMap2)
solve2 :: [String] -> Int64
solve2 = sumElements . foldl nextState (0, [], Map.empty)
where sumElements (_, _, map) = Map.foldl (+) 0 map
main :: IO ()
main = do
l <- lines <$> readFile "14.txt"
print $ solve2 l