adventofcode-2020/day13/main.hs

#! /usr/bin/env -S"GHCRTS=-N4" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple attoparsec arithmoi])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative
import Control.Monad (foldM)
import Data.Attoparsec.Text (Parser)
import Data.Euclidean (gcdExt)
import Data.List (find,sortOn)
import Data.Maybe (fromMaybe,catMaybes)
import Text.Pretty.Simple
import qualified Data.Attoparsec.Text as A
import qualified Data.Text as T

exampleData :: String
exampleData = "939\n7,13,x,x,59,x,31,19"

numOrXParser :: Parser (Maybe Int)
numOrXParser = (Just <$> A.decimal) <|> ("x" *> pure Nothing)

inputParser :: Parser (Int,[Int])
inputParser = do
  n <- A.decimal
  A.skipSpace
  xs <- numOrXParser `A.sepBy` ","
  pure (n,catMaybes $ xs)

parseInput :: String -> Either String (Int,[Int])
parseInput = (A.parseOnly inputParser) . T.pack

solvePart1 :: String -> Either String Int
solvePart1 str = do
  (n,xs) <- parseInput str
  let (bus, time) = head
                  $ sortOn (snd)
                  $ map (\x -> (x,fromMaybe (-1) $ find (> n) [0,x..])) xs
  pure $ bus * (time - n)

inputParser2 :: Parser (Int,[Maybe Int])
inputParser2 = do
  n <- A.decimal
  A.skipSpace
  xs <- numOrXParser `A.sepBy` ","
  pure (n,xs)

parseInput2 :: String -> Either String (Int,[Maybe Int])
parseInput2 = (A.parseOnly inputParser2) . T.pack

-- -------------------------------------------------------------------------- --
-- Here I sneak around a bit, and realize the problem is well defined         --
-- (and solved!) already: it's called the Chinese Remainder Theorem           --
-- https://en.wikipedia.org/wiki/Chinese_remainder_theorem                    --
-- -------------------------------------------------------------------------- --
-- Apllying the theorem allows to reduce a system of equation on x:           --
--                                                                            --
-- x ≡ a1 (mod n1)                                                            --
--   ·                                                                        --
--   ·                                                                        --
--   ·                                                                        --
-- x ≡ ak (mod nk)                                                            --
--                                                                            --
-- to a single equation:                                                      --
--                                                                            --
-- x ≡ as (mod ns)                                                            --
--                                                                            --
-- It relates to the buses schedules in the following way: t is x, the bus    --
-- number is the modulo factor (since a bus comes *every* ni) and subsequent  --
-- additions to t (for other buses) is (-ai), so, for a but coming at         --
-- t+ai, one would write x ≡ -ai (mod ni)                                     --
--                                                                            --
-- I chose to encode ai and ni as a tuple (ai,ni), named startAndIds          --
--                                                                            --
-- Basically, we're creating a “chinese” function:                            --
--                                                                            --
--   chinese :: (Int,Int) -> (Int,Int) -> (Int,Int)                           --
--                                                                            --
-- Then, given a list [(Int, Int)] we can fold over it to obtain the solution --
--                                                                            --

chinese :: (Integer,Integer) -> (Integer,Integer) -> Maybe (Integer,Integer)
chinese (0,n1) (0,n2)   = chinese (n1,n1) (n2,n2)
chinese v      (0,n2)   = chinese v       (n2,n2)
chinese (0,n1) v        = chinese (n1,n1) v
chinese (a1,n1) (a2,n2) = do
  -- Computes a solution such that: n1×c1 + n2×c2      = g, for some c2
  --                                n1×c1 - g          = - n2×c2, for some c2
  --                                1/n2 (n1×c1 - g)   = - c2, for some c2 (n2 is > 0)
  --                                - 1/n2 (n1×c1 - g) = c2, for some c2 (n2 is > 0)
  -- n1 and n2 must be coprimes for this to work (g must be 1), fail otherwise
  (m1,m2) <- case gcdExt n1 n2 of
             (1,c1) -> Just ( c1, negate ((n1 * c1) - 1) `div` n2 )
             _      -> Nothing
  let x = a1 * m2 * n2 + a2 * m1 * n1
  let a12 = x `mod` (n1 * n2)
  pure $ (a12, n1 * n2)

e2m :: Either e a -> Maybe a
e2m (Right v) = Just v
e2m _         = Nothing

solvePart2 :: String -> Maybe (Integer,Integer)
solvePart2 str = do
  (_,xs) <- e2m $ parseInput2 str
  let startAndIds = catMaybes $ sequence <$> zip [0..] (map (fmap fromIntegral) xs)
  let chineseEqs = fmap (\(a,n) -> ((-a) `mod` n, n)) startAndIds
  foldM chinese (1,1) chineseEqs

main :: IO ()
main = do
  putStrLn ":: Test"
  pPrint $ A.parseOnly inputParser $ T.pack exampleData
  pPrint $ take 3 ((\n -> [1,(n::Integer)..]) 59)
  putStrLn ":: Day 13 - Part 1"
  input <- readFile "day13/input"
  pPrint $ solvePart1 exampleData
  pPrint $ solvePart1 input
  putStrLn ":: Test 2"
  print $ solvePart2 exampleData
  print $ solvePart2 "1\n17,x,13,19"
  print $ solvePart2 "1\n67,7,59,61"
  print $ solvePart2 "1\n67,x,7,59,61"
  print $ solvePart2 "1\n67,7,x,59,61"
  print $ solvePart2 "1\n1789,37,47,1889"
  putStrLn ":: Day 13 - Part 2"
  print $ solvePart2 input