47 lines
1.6 KiB
Haskell
47 lines
1.6 KiB
Haskell
|
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
|