AoC/2021/day08/day08.hs

import Data.List.Split
import Data.List
import qualified Data.Map.Strict as M

main :: IO ()
main = do
  input <- map (map (map sort . words) . splitOn "|") <$> lines <$> readFile "day08/input"
  let output = map (!!1) input
  print $ length $ filter (`elem` [2,3,4,7]) $ map length (concat output)
  print $ day08 input

findDict :: [String] -> M.Map Char Char
findDict wire = M.fromList $ zip (a++b++c++d++e++f++g) "abcdefg"
    where a = intersect (getX 3) (without [2,4])
          b = intersect (charsWithLength 3 6) (charsWithLength 1 5)
          c = intersect (charsWithLength 2 6) (getX 2)
          d = intersect (charsWithLength 2 6) (charsWithLength 3 5)
          e = intersect (charsWithLength 1 5) $ intersect (getX 7) $ without [2..4]
          f = getX 2 \\ c
          g = without [2..4] \\ e
          without x = "abcdefg" \\ (concat $ filter (flip elem x . length) wire)
          getXs x = filter (\y->length y == x) wire
          getX = head . getXs
          charsWithLength x c = nub $ concat $ filter (\y->length y == x) $ group $ sort $ concat $ getXs c

decode :: M.Map Char Char -> String -> String
decode _ [] = []
decode dict (x:xs) = dict M.! x : decode dict xs

findNum :: String -> Int
findNum x
  | x=="abcefg" = 0
  | x=="cf" = 1
  | x=="acdeg" = 2
  | x=="acdfg" = 3
  | x=="bcdf" = 4
  | x=="abdfg" = 5
  | x=="abdefg" = 6
  | x=="acf" = 7
  | x=="abcdefg" = 8
  | x=="abcdfg" = 9

day08 :: [[[String]]] -> Int
day08 [] = 0
day08 ((wire:output):xs) = day08 xs + (sum $ zipWith (*) decoded [10^x|x<-[0..]])
    where dict = findDict wire
          decoded = reverse $ map (findNum . sort . decode dict) (head output)