Google Code Jam C - ThemePark

gistfile1.hs

module Main where

import Text.Printf
import Data.List

profit :: Int -> Int -> [Int] -> Int
profit cap rides groups 
  | sum groups <= cap = rides * sum groups
  | otherwise         = 
      case toCaps . findCycle $ rideCaps cap rides groups of
         (xs,[]) -> sum xs
         ([],ys) -> sumCycle rides ys
         (xs,ys) -> sum xs + sumCycle (rides - length xs) ys
  where
    toCaps (xs,ys) = (map fst xs, map fst ys)

sumCycle :: Int -> [Int] -> Int 
sumCycle rides cycle = a * sum cycle + sum (take b cycle)
    where
      l = length cycle
      a = rides `div` l
      b = rides `mod` l

-- the control
-- profit cap rides groups = sum . (map fst) $ rideCaps cap rides groups

-- returns stream of tuples with ride capacity and last group id in the ride
rideCaps :: Int -> Int -> [Int] -> [(Int,Int)]
rideCaps cap rides groups = rideCaps' cap rides (cycle (zip groups [0..])) 0
  where
    rideCaps' cap' rides' groups'@((g,c):gs) gcount'
      | rides' == 0 = []
      | cap' < g || gcount' == gcount = 
          (cap - cap', c) : rideCaps' cap (rides' - 1) groups' 0
      | otherwise = rideCaps' (cap' - g) rides' gs (gcount' + 1)
    gcount = length groups

processFile :: String -> String
processFile s = unlines $ zipWith (printf "Case #%d: %d") ([1..]::[Int]) profits
  where
    cs      = map (map read . words) . drop 1 . lines $ s
    profits = chop processCase cs

processCase :: [[Int]] -> (Int, [[Int]])
processCase ( [rides, cap, _ ] : groups : rest) = 
  (profit cap rides groups, rest)
processCase _ = error "Invalid Format!!!!!"

chop :: ([a] -> (b , [a])) -> [a] -> [b]
chop _ [] = []
chop f xs = y : chop f xs'
  where (y, xs') = f xs

main :: IO ()
main = interact processFile 

-- http://en.wikipedia.org/wiki/Floyd's_cycle-finding_algorithm
findCycle :: Eq a => [a] -> ([a],[a])
findCycle xxs = fCycle xxs xxs
  where fCycle _      []       = (xxs,[]) -- not cyclic
        fCycle _      [_]     =  (xxs,[])
        fCycle (x:xs) (_:y:ys)
         | x == y              = fStart xxs xs
         | otherwise           = fCycle xs ys
        fStart (x:xs) (y:ys)
         | x == y              = ([], x:fLength x xs)
         | otherwise           = let (as,bs) = fStart xs ys in (x:as,bs)
        fLength x (y:ys)
         | x == y              = []
         | otherwise           = y:fLength x ys