git.haldean.org aoc-2018 / master 04 / Aoc04.hs
master

Tree @master (Download .tar.gz)

Aoc04.hs @masterraw · history · blame

import Data.Fixed
import Data.Function
import Data.Maybe
import Data.Time
import Data.Time.Clock
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Set as S
import Text.Parsec
import Text.Parsec.String

data Record
  = Awaken UTCTime
  | Sleep UTCTime
  | ShiftStart UTCTime Int
  deriving (Show, Eq)


isAwaken :: Record -> Bool
isAwaken (Awaken _) = True
isAwaken _ = False

isSleep :: Record -> Bool
isSleep (Sleep _) = True
isSleep _ = False

isShiftStart :: Record -> Bool
isShiftStart (ShiftStart _ _) = True
isShiftStart _ = False

data RawShift = RawShift Int [Record] deriving Show
data Shift = Shift Int [(UTCTime, UTCTime)] deriving Show

recordTime :: Record -> UTCTime
recordTime (Awaken t) = t
recordTime (Sleep t) = t
recordTime (ShiftStart t _) = t

recordDay :: Record -> Day
recordDay = utctDay . recordTime

buildRecordTime :: Int -> Int -> Int -> Int -> Int -> Maybe UTCTime
buildRecordTime y mo d h m =
  fromGregorianValid (toInteger y) mo d
  >>= (\day -> (Just $ UTCTime day (secondsToDiffTime . toInteger $ 60 * (60 * h + m))))
  -- add 1hr to time so that the shift records all sit in the same day
  >>= Just . addUTCTime (fromInteger 3600)

instance Ord Record where
  compare = compare `on` recordTime

parseInt :: Parser Int
parseInt = fmap (read :: String -> Int) $ many1 digit

parseGuard :: Parser Int
parseGuard = do
  _ <- string "Guard #"
  gid <- parseInt
  _ <- string " begins shift"
  return gid

parseAwaken :: Parser ()
parseAwaken = do
  _ <- string "wakes up"
  return ()

parseSleep :: Parser ()
parseSleep = do
  _ <- string "falls asleep"
  return ()

parseLine :: Parser Record
parseLine = do
  _ <- char '['
  year <- parseInt
  _ <- char '-'
  month <- parseInt
  _ <- char '-'
  day <- parseInt
  _ <- char ' '
  hour <- parseInt
  _ <- char ':'
  minute <- parseInt
  _ <- many $ oneOf "] "
  case buildRecordTime year month day hour minute of
    Just time -> (parseGuard >>= \gid -> return $ ShiftStart time gid)
                  <|> (parseAwaken >> (return $ Awaken time))
                  <|> (parseSleep >> (return $ Sleep time))
    Nothing -> undefined

partitionShifts :: [Record] -> [[Record]]
partitionShifts = L.groupBy ((==) `on` recordDay) . L.sort

findGuard :: [Record] -> Maybe RawShift
findGuard shiftRecords =
  let guardRecords = filter isJust $ map getGuardId shiftRecords where
        getGuardId (ShiftStart _ gid) = Just gid
        getGuardId _ = Nothing
  in case guardRecords of
    [(Just gid)] -> Just $ RawShift gid $ filter (not . isShiftStart) shiftRecords
    _ -> error $ "bad set of records for shift: " ++ show shiftRecords

pairUpTimes :: RawShift -> Shift
pairUpTimes (RawShift gid records) =
  let sleeps = map recordTime $ filter isSleep records
      awakens = map recordTime $ filter isAwaken records
  in Shift gid $ zip sleeps awakens

getShifts :: [Record] -> [Shift]
getShifts = map pairUpTimes . catMaybes . map findGuard . partitionShifts

getSleepTime :: Shift -> Integer
getSleepTime (Shift _ pairs) = sum $ map (\(t1, t2) -> floor $ (diffUTCTime t2 t1) / 60) pairs

shiftGuard :: Shift -> Int
shiftGuard (Shift gid _) = gid

allGuards :: [Shift] -> S.Set Int
allGuards = S.fromList . map shiftGuard

guardSleepTimes :: [Shift] -> [(Integer, Int)] -- sleep time, guard id
guardSleepTimes shifts =
  shifts
  & allGuards
  & S.toList
  & map (\gid -> (sum $ map getSleepTime $ filter (\s -> shiftGuard s == gid) shifts, gid))
  & L.sort

sleepiestGuard :: [Shift] -> Int
sleepiestGuard = snd . L.maximum . guardSleepTimes

minute :: UTCTime -> Int
minute = (\s -> floor $ (s `mod'` 3600) / 60) . utctDayTime

isAsleepAt :: Int -> Shift -> Bool
isAsleepAt m (Shift _ pairs) = any id $ map minuteIsBetween pairs where
  minuteIsBetween (t1, t2) = (minute t1 <= m) && (m < minute t2)

guardMinutes :: Int -> [Shift] -> [(Int, Int)]
guardMinutes g allShifts =
  let shifts = filter ((==) g . shiftGuard) allShifts
  in map (countSleepMinutes shifts) [0..59] where
    countSleepMinutes shifts m = (length $ filter (isAsleepAt m) shifts, m)

sleepiestMinute :: Int -> [Shift] -> (Int, Int) -- sleeps, minute
sleepiestMinute g allShifts = L.maximum $ guardMinutes g allShifts

sleepiestGuardMinute :: [Shift] -> (Int, Int, Int) -- sleeps, minute, guard
sleepiestGuardMinute allShifts =
  let guards = S.toList $ allGuards allShifts
  in L.maximum $ map (\g -> let (s, m) = sleepiestMinute g allShifts in (s, m, g)) guards

main :: IO ()
main = do
  records <- fmap (parse (sepEndBy parseLine endOfLine) "input.txt") (readFile "input.txt")
  case records of
    Right rs -> do
      let shifts = getShifts rs
      let guard = sleepiestGuard shifts
      let guardMin = snd $ sleepiestMinute guard shifts
      putStrLn $ "1: sleepiest guard:  " ++ show guard
      putStrLn $ "1: sleepiest minute: " ++ show guardMin
      putStrLn $ "1: result: " ++ show (guard * guardMin)
      let (s2, m2, g2) = sleepiestGuardMinute shifts
      putStrLn $ "2: sleepiest guard:  " ++ show g2
      putStrLn $ "2: sleepiest minute: " ++ show m2
      putStrLn $ "2: result: " ++ show (g2 * m2)
    Left err -> print err