blob: 4a4e24dac73e0b030abb5bd03b73f4d7aa088ff7 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
|
module Main where
import Data.Bits
import qualified Data.IntMap.Strict as IM
import Data.IntMap.Strict (IntMap)
import Data.List (foldl')
import Input
subsets :: [a] -> [[a]]
subsets [] = [[]]
subsets (x:xs) = subsets xs ++ map (x:) (subsets xs)
data Mask = Mask Int Int
deriving (Show)
applyMask :: Mask -> Int -> Int
applyMask (Mask set reset) n = (n .|. set) .&. complement reset
data Instr = SetMask Mask | Store Int Int
deriving (Show)
parseInstr :: String -> Instr
parseInstr str
| str !! 1 == 'a' = SetMask (parseMask (drop 7 str))
| [(idx, rest)] <- reads (drop 4 str) = Store idx (read (drop 4 rest))
| otherwise = error "Can't parse Instr"
parseMask :: String -> Mask
parseMask s =
let pairs = zip s (map bit [35,34..0])
setbits = sum . map snd . filter ((== '1') . fst) $ pairs
resetbits = sum . map snd . filter ((== '0') . fst) $ pairs
in Mask setbits resetbits
data State = State { sMask :: Mask, sMem :: IntMap Int }
deriving (Show)
exec1 :: State -> Instr -> State
exec1 (State _ mem) (SetMask mask) = State mask mem
exec1 (State mask mem) (Store idx val) =
State mask (IM.insert idx (applyMask mask val) mem)
expandWithMask :: Mask -> Int -> [Int]
expandWithMask (Mask ones zeros) n =
let floating = (bit 36 - 1) .&. complement (ones .|. zeros)
n' = (n .|. ones) .&. complement floating
singlebits = [bit i | i <- [0..35], testBit floating i]
in map ((n' .|.) . sum) (subsets singlebits)
exec2 :: State -> Instr -> State
exec2 (State _ mem) (SetMask mask) = State mask mem
exec2 (State mask mem) (Store idx val) =
State mask (foldr (\idx' -> IM.insert idx' val) mem (expandWithMask mask idx))
main :: IO ()
main = do
input <- map parseInstr <$> getInput 14
print (sum . IM.elems . sMem $ foldl' exec1 (State (Mask 0 0) mempty) input)
print (sum . IM.elems . sMem $ foldl' exec2 (State (Mask 0 0) mempty) input)
|
