view optim.py @ 15:d9aa5e25859c draft

main: use CLI args to select the algorithm to use Also return maxplan from brute_force() so it can be used in main() with the same API as greedy().
author Jordi Gutiérrez Hermoso <jordigh@octave.org>
date Wed, 11 Mar 2015 16:47:17 -0400
parents a72cf3c7c976
children 4c23cfdeaa5f
line wrap: on
line source

#!/usr/bin/env python

from collections import namedtuple

Machine = namedtuple("Machine", ["day", "buy", "sell", "profit", "maxprofit"])
Case = namedtuple("Case", ["machines", "days", "capital"])

from argparse import ArgumentParser

def process_options():
    parser = ArgumentParser(
        description=("Solve problem F."))
    parser.add_argument("file", help = "input file to read")
    parser.add_argument("--debug", dest = "debug", action = "store_true",
                        help = "enable debug output")

    algorithm = parser.add_mutually_exclusive_group(required = False)

    algorithm.add_argument("-b", "--brute-force", dest = "brute_force",
                           action = "store_true",
                           help = "Use brute force, O(2**N) time")
    algorithm.add_argument("-g", "--greedy", dest = "greedy",
                           action = "store_true", default = True,
                           help = ("Use greedy algorithm, fails for some cases, "
                                   "roughly O(N) time"))

    args = parser.parse_args()
    return args

def parseinput(fname):
    """
    Parse the input file, forget about input validation
    """
    cases = []
    with open(fname) as f:
        while True:
            header = [int(x) for x in f.readline().split()]
            if header == [0, 0, 0]:
                return cases
            N = header[0]
            case = Case([], header[1], header[2])
            for i in range(0, N):
                machine = Machine(*[int(x) for x in f.readline().split()],
                                  maxprofit = None)

                # Maximum profit possible from each machine
                maxprofit = ((case.days - machine.day)*machine.profit
                             - machine.buy + machine.sell)

                # Skip machines that can only give a loss. This
                # includes machines that can only be bought after
                # case.days.
                if maxprofit < 0:
                    continue

                machine = machine._replace(maxprofit = maxprofit)
                case.machines.append(machine)

            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 first sell the old machine, if any
            if slot:
                if machine.day == slot.day:
                    raise InvalidPlan(
                        "Cannot buy two machines on the same day: %d"
                        % (machine.day))

                # 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

            # Then we try to buy the new one
            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
            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 brute_force(case):
    """Form all possible plans, and just try them all."""
    machines = case.machines

    # Just count in binary to enumerate all plans
    N = len(machines)
    if N > 0:
        plans = [[y == "1" for y in format(x, "0%db" % N)]
                 for x in range(0, 2**N)]
    else:
        plans = []

    maxcapital = case.capital
    maxplan = None
    for plan in plans:
        try:
            plancapital = final_capital(case, machines, plan)
            if plancapital > maxcapital:
                maxcapital = plancapital
                maxplan = plan
        except InvalidPlan:
            pass

    return maxcapital, maxplan

def greedy(case):
    """Greedy algorithm, considering each machine only once, in decreasing
    order of maxprofit.

    """
    machines = sorted(case.machines, key = lambda m: -m.maxprofit)
    N = len(machines)
    maxcapital = case.capital

    # Base plan, pick no machine
    plan = [False]*len(machines)
    maxplan = plan

    for idx in range(0, N):
        plan[idx] = True
        try:
            plancapital = final_capital(case, machines, plan)
            if plancapital > maxcapital:
                maxcapital = plancapital
                maxplan = plan
            else:
                plan[idx] = False
        except InvalidPlan:
            plan[idx] = False

    return maxcapital, maxplan

def main():
    args = process_options()
    cases = parseinput(args.file)
    for (number, case) in enumerate(cases):
        if args.greedy:
            solver = greedy
        if args.brute_force:
            solver = brute_force
        maxcapital, plan = solver(case)
        print "Case %d: %d" % (number + 1, maxcapital)

if __name__ == "__main__":
    main()