advent08/Main-longwinded.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/09/advent-of-code-2021-day-8/
   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 hiding ((\\))
  10 import qualified Data.Map.Strict as M
  11 import Data.Map.Strict ((!))
  12 import qualified Data.Set as S
  13 import Data.Set ((\\))
  14
  15 data Display = Display [String] [String] -- patterns, output
  16   deriving (Eq, Show)
  17
  18 type Assignments = M.Map Segment (S.Set Char)
  19 type Encoding = M.Map Char Segment
  20 type DigitSegments = M.Map (S.Set Segment) Char
  21
  22 data Segment = Seg1 | Seg2 | Seg3 | Seg4 | Seg5 | Seg6 | Seg7
  23   deriving (Eq, Ord, Show)
  24
  25
  26 main :: IO ()
  27 main =
  28   do  text <- TIO.readFile "data/advent08.txt"
  29       let displays = successfulParse text
  30       print $ part1 displays
  31       print $ part2 displays
  32       -- let td = successfulParse "acedgfb cdfbe gcdfa fbcad dab cefabd cdfgeb eafb cagedb ab | cdfeb fcadb cdfeb cdbaf"
  33       -- print $ part2 td
  34
  35 part1 displays = sum $ map countUniques displays
  36
  37 countUniques (Display _ outputs) = length uniqLens
  38   where outLens = map length outputs
  39         uniqLens = outLens `intersect` uniqueLengths
  40
  41 uniqueLengths = [2, 3, 4, 7]
  42
  43 segmentNames = S.fromList "abcdefg"
  44
  45 digitSegments :: DigitSegments
  46 digitSegments = M.fromList
  47   [ (S.fromList [Seg1, Seg2, Seg3, Seg5, Seg6, Seg7], '0')
  48   , (S.fromList [Seg3, Seg6], '1')
  49   , (S.fromList [Seg1, Seg3, Seg4, Seg5, Seg7], '2')
  50   , (S.fromList [Seg1, Seg3, Seg4, Seg6, Seg7], '3')
  51   , (S.fromList [Seg2, Seg3, Seg4, Seg6], '4')
  52   , (S.fromList [Seg1, Seg2, Seg4, Seg6, Seg7], '5')
  53   , (S.fromList [Seg1, Seg2, Seg4, Seg5, Seg6, Seg7], '6')
  54   , (S.fromList [Seg1, Seg3, Seg6], '7')
  55   , (S.fromList [Seg1, Seg2, Seg3, Seg4, Seg5, Seg6, Seg7], '8')
  56   , (S.fromList [Seg1, Seg2, Seg3, Seg4, Seg6, Seg7], '9')
  57   ]
  58
  59
  60 -- combine = M.unionWith intersect
  61
  62 part2 displays = sum $ map decodeOneDisplay displays
  63
  64 decodeOneDisplay display = findCode invAllocation display
  65   where allocation = allocate display
  66         invAllocation = invertAssignment allocation
  67
  68 allocate :: Display -> Assignments
  69 allocate (Display examples _) = assignments6
  70   where
  71     -- segments 3 and 6 are given by pattern of length 2 (digit 1)
  72     dSegs1 = S.fromList $ head $ filter ((== 2) . length) examples
  73     assignments0 = M.fromList [(Seg3, dSegs1), (Seg6, dSegs1)]
  74     -- segment 1 is the one in pattern of length 3 (digit 7) that's
  75     -- not in segments for digit 1
  76     dSegs7 = S.fromList $ head $ filter ((== 3) . length) examples
  77     assignments1 = M.insert Seg1 (dSegs7 \\ dSegs1) assignments0
  78     -- segments 2 and 4 are the ones in pattern of length 4 (digit 4)
  79     -- that aren't in pattern for digit 1
  80     dSegs4 = (S.fromList $ head $ filter ((== 4) . length) examples) \\ dSegs1
  81     -- segments 1, 4, 7 are the common ones in digits 2, 3, 5
  82     dSegs235 = S.fromList $ foldl1' intersect $ filter ((== 5) . length) examples
  83     -- segment 4 is the common one with digit 4
  84     seg4 = dSegs235 `S.intersection` dSegs4
  85     -- segment 2 is the other one from digit 4
  86     seg2 = dSegs4 \\ seg4
  87     assignments2 = M.union (M.fromList [(Seg2, seg2), (Seg4, seg4)]) assignments1
  88     -- we now know segments 1 and 4, so deduce segment 7
  89     seg7 = dSegs235 \\ (S.union (assignments2!Seg1) (assignments2!Seg4))
  90     assignments3 = M.union (M.singleton Seg7 seg7) assignments2
  91     -- of the 5-segment digits, segment 2 only in digit 5
  92     seg2c = head $ S.toList seg2
  93     dSegs5 = S.fromList $ head $ filter (elem seg2c) $ filter ((== 5) . length) examples
  94     -- remove known values of segments 1, 2, 4, 7
  95     segs1247 = S.unions $ M.elems $ assignments3 `M.restrictKeys` (S.fromList [Seg1, Seg2, Seg4, Seg7])
  96     -- segs1247 = S.unions [assignments3!Seg1, assignments3!Seg2,
  97     --                       assignments3!Seg4, assignments3!Seg7]
  98     -- what's left is segment 6
  99     seg6 = dSegs5 \\ segs1247
 100     assignments4 = M.insert Seg6 seg6 assignments3
 101     -- segment 3 can't be the same allocation as segment 6
 102     assignments5 = M.insert Seg3 ((assignments4!Seg3) \\ seg6) assignments4
 103     -- segment 5 is the only one left
 104     seg5 = segmentNames \\ (S.unions $ M.elems assignments5)
 105     assignments6 = M.insert Seg5 seg5 assignments5
 106
 107
 108 invertAssignment assignment = M.foldrWithKey inv1 M.empty assignment
 109   where inv1 seg vals invMap = M.insert (head $ S.elems vals) seg invMap
 110
 111
 112 findDigit :: Encoding -> [Char] -> Char
 113 findDigit segmentAssignments code = digitSegments ! segments
 114   where codeSet = S.fromList code
 115         segmentMap = M.restrictKeys segmentAssignments codeSet
 116         segments = S.fromList $ M.elems segmentMap
 117
 118 findDigits :: Encoding -> [[Char]] -> [Char]
 119 findDigits segmentAssignments codes = map (findDigit segmentAssignments) codes
 120
 121 findCode :: Encoding -> Display -> Int
 122 findCode segmentAssignments (Display _ codes) = read $ findDigits segmentAssignments codes
 123
 124
 125
 126 -- Parse the input file
 127
 128 displaysP = displayP `sepBy` endOfLine
 129 displayP = Display <$> (patternsP <* " | ") <*> patternsP
 130
 131 patternsP = patternP `sepBy` " "
 132 patternP = many1 letter
 133
 134 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 135 successfulParse input =
 136   case parseOnly displaysP input of
 137     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 138     Right displays -> displays