advent08/Main.hs

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