Skip to content

Instantly share code, notes, and snippets.

@petermarks
Created January 22, 2011 17:27
Show Gist options
  • Save petermarks/791274 to your computer and use it in GitHub Desktop.
Save petermarks/791274 to your computer and use it in GitHub Desktop.
Theme Park (Optimized)
-- 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment