advent08/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/04/advent-of-code-2021-day-4/
   2
   3 import Data.Text ()
   4 import qualified Data.Text.IO as TIO
   5
   6 import Data.Attoparsec.Text
   7 -- import Control.Applicative
   8
   9 import Data.List
  10 import qualified Data.Map.Strict as M
  11 import Data.Map.Strict ((!))
  12 import qualified Data.Set as S
  13
  14 data Display = Display [String] [String] -- patterns, output
  15   deriving (Eq, Show)
  16
  17 type Encoding = M.Map Char Segment
  18 type DigitSegments = M.Map (S.Set Segment) Char
  19
  20 data Segment = Seg1 | Seg2 | Seg3 | Seg4 | Seg5 | Seg6 | Seg7
  21   deriving (Eq, Ord, Show, Enum, Bounded)
  22
  23
  24 main :: IO ()
  25 main =
  26   do  text <- TIO.readFile "data/advent08.txt"
  27       let displays = successfulParse text
  28       print $ part1 displays
  29       print $ part2 displays
  30
  31 part1 :: [Display] -> Int
  32 part1 displays = sum $ map countUniques displays
  33
  34 part2 :: [Display] -> Int
  35 part2 displays = sum $ map decodeOneDisplay displays
  36
  37 countUniques :: Display -> Int
  38 countUniques (Display _ outputs) = length uniqLens
  39   where outLens = map length outputs
  40         uniqLens = outLens `intersect` uniqueLengths
  41
  42 uniqueLengths :: [Int]
  43 uniqueLengths = [2, 3, 4, 7]
  44
  45 decodeOneDisplay :: Display -> Int
  46 decodeOneDisplay display = findCode allocation display
  47   where allocation = allocate display
  48
  49 segmentNames :: [Char]
  50 segmentNames = "abcdefg"
  51
  52 segments :: [Segment]
  53 segments = [Seg1 .. Seg7]
  54
  55 digitSegments :: DigitSegments
  56 digitSegments = M.fromList
  57   [ (S.fromList [Seg1, Seg2, Seg3, Seg5, Seg6, Seg7], '0')
  58   , (S.fromList [Seg3, Seg6], '1')
  59   , (S.fromList [Seg1, Seg3, Seg4, Seg5, Seg7], '2')
  60   , (S.fromList [Seg1, Seg3, Seg4, Seg6, Seg7], '3')
  61   , (S.fromList [Seg2, Seg3, Seg4, Seg6], '4')
  62   , (S.fromList [Seg1, Seg2, Seg4, Seg6, Seg7], '5')
  63   , (S.fromList [Seg1, Seg2, Seg4, Seg5, Seg6, Seg7], '6')
  64   , (S.fromList [Seg1, Seg3, Seg6], '7')
  65   , (S.fromList [Seg1, Seg2, Seg3, Seg4, Seg5, Seg6, Seg7], '8')
  66   , (S.fromList [Seg1, Seg2, Seg3, Seg4, Seg6, Seg7], '9')
  67   ]
  68
  69 allocate :: Display -> Encoding
  70 allocate (Display examples _) = head $ validEncodings
  71   where allEncodings = map (\segs -> M.fromList $ zip segmentNames segs)
  72                             $ permutations segments
  73         validEncodings = filter (isValidEncoding examples) allEncodings
  74
  75 segmentsOfSignal :: Encoding -> [Char] -> S.Set Segment
  76 segmentsOfSignal encoding signal = S.fromList $ map (encoding ! ) signal
  77
  78 isValidEncoding :: [[Char]] -> Encoding -> Bool
  79 isValidEncoding examples encoding =
  80   all (\e -> M.member e digitSegments) exampleSegments
  81   where exampleSegments = map (segmentsOfSignal encoding) examples
  82
  83
  84 findDigit :: Encoding -> [Char] -> Char
  85 findDigit encoding code = digitSegments ! (segmentsOfSignal encoding code)
  86
  87 findDigits :: Encoding -> [[Char]] -> [Char]
  88 findDigits encoding codes = map (findDigit encoding) codes
  89
  90 findCode :: Encoding -> Display -> Int
  91 findCode encoding (Display _ codes) = read $ findDigits encoding codes
  92
  93
  94
  95 -- Parse the input file
  96
  97 displaysP = displayP `sepBy` endOfLine
  98 displayP = Display <$> (patternsP <* " | ") <*> patternsP
  99
 100 patternsP = patternP `sepBy` " "
 101 patternP = many1 letter
 102
 103 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 104 successfulParse input =
 105   case parseOnly displaysP input of
 106     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 107     Right displays -> displays