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