Theme Park (Optimized)

gistfile1.hs

-- This program solves the Theme Park Google CodeJam problem at
-- http://code.google.com/codejam/contest/dashboard?c=433101#s=p2
--
-- I've tried to make the code clear whilst employing a selection of Haskell library
-- functions and idioms. I've also gone for efficiency, probably over optimizing in
-- places in order to demonstrate some techniques.

module Main where

import Data.List
import Data.Array
import Text.Printf

-- The bigest number we need to handle is 10^8 * 10^9 which won't fit in 32 bits,
-- but will in 64. We could use Word64, but Integer will work just fine.

type N = Integer


-- Calculate the profit for a scenario.
--
-- To perform acceptably, we need to detect cycles in the list ride loading.

profit :: N -> N -> [N] -> N
profit cap rides groups 
  | sum groups <= cap = rides * sum groups
  | otherwise         = initialProfit + cycleProfit + finalProfit
  where
    initialProfit          = rideProfit cycleStart
    cycleProfit            = (rideProfit cycleEnd - rideProfit cycleStart) * numCycles
    numCycles              = (rides - rideCount cycleStart) `div` cycleLength
    finalProfit            = rideProfit finalRide - rideProfit cycleEnd
    (cycleStart, cycleEnd) = findCycle $ calcRides cap rides groups
    finalRide              = genericDrop numRidesAfterCycle cycleEnd
    numRidesAfterCycle     = rides - cycleLength * numCycles - rideCount cycleStart
    cycleLength            = rideCount cycleEnd - rideCount cycleStart


-- A Ride holds the state after each ride.
-- All members are strict so that accumulators accumulate rather than creating thunk chains.

data Ride = Ride
  { rCount     :: !N   -- The number of rides to this point
  , rProfit    :: !N   -- The profit so far
  , rNextIndex :: !Int -- The position in the original queue of the next group to board
  }
  deriving (Show)


-- A Loading represents a loading of the rollercoaster. 

data Loading = Loading
  { lProfit     :: !N   -- The profit for this loading
  , lNextIndex  :: !Int -- The position in the original queue of the next group to board
  }


-- Generate a list of rides for a scenario

calcRides :: N -> N -> [N] -> [Ride]
calcRides cap rides groups = iterate nextRide (Ride 0 0 0)
  where
    nextRide (Ride c p i) = let (Loading p' i') = loadings ! i in Ride (c + 1) (p + p') i'
    loadings = calcLoadings cap groups


-- Generate an array of loadings from each possible starting position in the queue.
--
-- This is a memoization of the load function below. As we are interested in indeces later,
-- we convert the list of groups into an array and just track indices. Whilst we could call
-- load with an empty rollercoaster at each queue position, we optimize by loading the
-- coasater once then iteratively removing one group and filling just the empty seats.
-- The state for our unfold is the previous Loading and the list of groups to unload on each
-- future iteration.

calcLoadings :: N -> [N] -> Array Int Loading
calcLoadings cap groups = listArray (0, length groups - 1) . unfoldr nextLoading $ (Loading 0 0, 0 :  groups)
  where
    nextLoading ((Loading p i), g:gs) = 
      let l = load cap groupsArray (Loading (p - g) i) in Just (l, (l, gs))
    groupsArray = listArray (0, length groups - 1) groups


-- Complete the loading of a (possibly) partially loaded coaster.

load :: N -> Array Int N -> Loading -> Loading
load cap groups = until full nextGroup
  where
    full      (Loading p i) = p + groups ! i > cap
    nextGroup (Loading p i) = Loading (p + groups ! i) ((i + 1) `mod` numGroups)
    numGroups               = snd (bounds groups) + 1


-- Some utility functions for accessing information about the next ride in a list of rides.

rideCount :: [Ride] -> N
rideCount = rCount . head

rideProfit :: [Ride] -> N
rideProfit = rProfit . head

rideNextIndex :: [Ride] -> Int
rideNextIndex = rNextIndex . head


-- Find a cycle in the rides list.
--
-- We use a cut down version of Floyd's cycle-finding algorithm:
-- http://en.wikipedia.org/wiki/Cycle_detection. We don't need to find the smallest cycle,
-- just a cycle, so as soon as the tortoise and the hare are at the same index value, the
-- position of the tortoise is at the start of a cycle and the difference in position of
-- the hare and the tortoise is the length of the cycle.

findCycle :: [Ride] -> ([Ride], [Ride])
findCycle rides = until sameIndex move (rides, tail rides)
  where
    sameIndex (tortoise, hare) = rideNextIndex tortoise == rideNextIndex hare
    move      (tortoise, hare) = (tail tortoise, tail . tail $ hare)


-- Parse the input, process each test case and generate the output.
    
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


-- Process an individual test case.

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


-- Lennart's chop function

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


-- Main just calls processFile with standard in and directs the output to standard out.

main :: IO ()
main = interact processFile