aoc2021/app/D3P2.hs

module Main where

import Data.Bits (shift)


main :: IO ()
main = interact solution

solution :: String -> String
solution = show
  . (\(x,y) -> bits_to_number x * bits_to_number y)
  . (\x -> (o2 0 x, co2 0 x))
  . map (map (read . (:[])))
  . lines
  where
    bits_to_number :: [Int] -> Int
    bits_to_number = btn' . reverse
      where
        btn' :: [Int] -> Int
        btn' []    = 0
        btn' (a:r) = shift (btn' r) 1 + a

count :: Eq a => Int -> a -> [[a]] -> Int
count _ _ []    = 0
count p e (a:r) = count p e r + if a !! p == e then 1 else 0


o2 :: Int -> [[Int]] -> [Int]
o2 _ [n] = n
o2 n x   = o2 next (filter (\e -> e !! n == mc) x)
  where
    next = n + 1
    c0   = count n 0 x
    c1   = count n 1 x
    mc
      | c0 <= c1 = 0
      | otherwise = 1

co2 :: Int -> [[Int]] -> [Int]
co2 _ [n] = n
co2 n x   = co2 next (filter (\e -> e !! n == lc) x)
  where
    next = n + 1
    c0   = count n 0 x
    c1   = count n 1 x
    lc
      | c0 <= c1 = 1
      | otherwise = 0