Mercurial > hg > problem6
diff optim.py @ 4:14c8b6dad88a draft
final_capital: new function
This is the objective function that we must maximise with a buying plan
author | Jordi GutiƩrrez Hermoso <jordigh@octave.org> |
---|---|
date | Tue, 10 Mar 2015 21:21:03 -0400 (2015-03-11) |
parents | c683f80bc858 |
children | cd24a8027d5f |
line wrap: on
line diff
--- a/optim.py +++ b/optim.py @@ -28,6 +28,53 @@ cases.append(case) +class InvalidPlan(Exception): + pass + +def final_capital(case, machines, plan): + """Given a case, some machines, and an action plan about whether to + buy a machine or not, compute the final capital given by this + plan. Raises InvalidPlan if it's impossible to buy a machine at + any step. + + This is our objective function to maximise over all possible + plans. + + """ + capital = case.capital + + # The "slot" that holds a machine. Initially, there's nothing + # there. + slot = None + + for (action, machine) in sorted(zip(plan, machines), + key = lambda x: x[1].day): + if action: + # We attempt to buy a new machine + if capital < machine.buy: + raise InvalidPlan("Not enough capital at day %d: " + "machine cost is %d, and capital is %d" + % (machine.day, machine.buy, capital)) + + capital -= machine.buy + if slot: + # Subtract 1, because the machine in the slot cannot + # be used the day it's sold. + days_used = machine.day - slot.day - 1 + capital += days_used*slot.profit + capital += slot.sell + + slot = machine + + # We account for the final machine, if we have it. + if slot: + # We can sell the day after the restructuring period, so no + # subtracting 1 here. + days_used = case.days - slot.day + capital += days_used*slot.profit + capital += slot.sell + + return capital def main(): cases = parseinput("input.txt") for case in cases: