solve.py

def solve(duties, trips):
  problem = LpProblem('driver_scheduling', LpMinimize)
  variables = []
  costs = []
  # Data structure to generate constraints for each trip.
  variables_for_trip = {trip: [] for trip in trips}

  # We have to make sure there's some set of duties that really do include all
  # the trips. Since we randomly generate duties, we can't be sure. E.g., we
  # could get unlucky and randomly generate duties that all only have trips on
  # the 'A' route. To solve this problem, generate one duty for each trip that
  # includes only that trip.
  duties = duties + [[trip] for trip in trips]

  # Gather up variables and costs
  for i, duty in enumerate(duties):
    # Create a binary variable (the value can be only 0 or 1)
    x = LpVariable('x{}'.format(i + 1), 0, 1, LpBinary)
    variables.append(x)
    costs.append(cost(duty))
    for trip in duty:
      variables_for_trip[trip].append(x)

  # Create the objective function. lpDot is shorthand for
  # c * x for (c, x) in zip(costs, variables)
  problem += lpDot(costs, variables)

  # Create constraints that for each trip, exactly one x from the duties
  # including it must be 1.
  for xs in variables_for_trip.values():
    problem += lpSum(xs) == 1

  # Pulp gives a very nice string representation of the problem when printed.
  print(problem)
  status = problem.solve()
  print(LpStatus[status])

  # We have a solution, now look at the values of xs to determine which duties
  # to use. Sum the cost for each used duty.
  solution = []
  total_cost = 0
  for i, x in enumerate(variables):
    if x.value():
      solution.append(duties[i])
      total_cost += costs[i]

  return solution, total_cost