Assignment with Teams of Workers

There are many versions of the assignment problem, which have additional constraints on the workers or tasks. In the next example, six workers are divided into two teams, and each team can perform at most two tasks.

The following sections present a Python program that solves this problem using the MIP solver. For a solution using the min cost flow solver, see the section Assignment with teams.

Declare the solver

The following code declares the MIP solver.

from ortools.linear_solver import pywraplp

def main():
  solver = pywraplp.Solver('SolveAssignmentProblemMIP',
                           pywraplp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)

Create the data

The following code sets up the data for the problem.

  cost = [[90, 76, 75, 70],
          [35, 85, 55, 65],
          [125, 95, 90, 105],
          [45, 110, 95, 115],
          [60, 105, 80, 75],
          [45, 65, 110, 95]]

  team1 = [0, 2, 4]
  team2 = [1, 3, 5]
  team_max = 2
The array cost is the cost matrix.

Create the variables

The following code creates binary integer variables for the problem.

  x = {}

  for i in range(num_workers):
    for j in range(num_tasks):
      x[i, j] = solver.BoolVar('x[%i,%i]' % (i, j))

There is one variable for each pair of a worker and task. Note that the workers are numbered 0 — 5, while the tasks are numbered 0 — 3, unlike in the original example, in which all nodes had to be numbered differently, as required by the min cost flow solver.

Create the objective function

The following code creates the objective function for the problem.

  solver.Minimize(solver.Sum([cost[i][j] * x[i,j] for i in range(num_workers)
                                                  for j in range(num_tasks)]))

The value of the objective function is the total cost over all variables that are assigned the value 1 by the solver.

Create the constraints

The following code creates the constraints for the problem.


  # Each worker is assigned to at most 1 task.

  for i in range(num_workers):
    solver.Add(solver.Sum([x[i, j] for j in range(num_tasks)]) <= 1)

  # Each task is assigned to exactly one worker.

  for j in range(num_tasks):
    solver.Add(solver.Sum([x[i, j] for i in range(num_workers)]) == 1)

  # Each team takes on two tasks.

  solver.Add(solver.Sum([x[i, j] for i in team1 for j in range(num_tasks)]) <= team_max)
  solver.Add(solver.Sum([x[i, j] for i in team2 for j in range(num_tasks)]) <= team_max)

Invoke the solver

The following code invokes the solver for the problem.

  sol = solver.Solve()

  print('Total cost = ', solver.Objective().Value())
  print()
  for i in range(num_workers):
    for j in range(num_tasks):
      if x[i, j].solution_value() > 0:
        print('Worker %d assigned to task %d.  Cost = %d' % (
              i,
              j,
              cost[i][j]))

  print()
  print("Time = ", solver.WallTime(), " milliseconds")

The command

print("Time = ", solver.WallTime(), " milliseconds")

displays the time the program takes to run.

Here is the output of the program.

Minimum cost assignment:  250.0

Worker 0 assigned to task 2.  Cost = 75
Worker 1 assigned to task 0.  Cost = 35
Worker 4 assigned to task 3.  Cost = 75
Worker 5 assigned to task 1.  Cost = 65

Time =  6  milliseconds

Entire program

Here is the entire code for the program.

from __future__ import print_function
from ortools.linear_solver import pywraplp

def main():
  solver = pywraplp.Solver('SolveAssignmentProblemMIP',
                           pywraplp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)

  cost = [[90, 76, 75, 70],
          [35, 85, 55, 65],
          [125, 95, 90, 105],
          [45, 110, 95, 115],
          [60, 105, 80, 75],
          [45, 65, 110, 95]]

  team1 = [0, 2, 4]
  team2 = [1, 3, 5]
  team_max = 2

  num_workers = len(cost)
  num_tasks = len(cost[1])
  x = {}

  for i in range(num_workers):
    for j in range(num_tasks):
      x[i, j] = solver.BoolVar('x[%i,%i]' % (i, j))

  # Objective
  solver.Minimize(solver.Sum([cost[i][j] * x[i,j] for i in range(num_workers)
                                                  for j in range(num_tasks)]))

  # Constraints

  # Each worker is assigned to at most 1 task.

  for i in range(num_workers):
    solver.Add(solver.Sum([x[i, j] for j in range(num_tasks)]) <= 1)

  # Each task is assigned to exactly one worker.

  for j in range(num_tasks):
    solver.Add(solver.Sum([x[i, j] for i in range(num_workers)]) == 1)

  # Each team takes on two tasks.

  solver.Add(solver.Sum([x[i, j] for i in team1 for j in range(num_tasks)]) <= team_max)
  solver.Add(solver.Sum([x[i, j] for i in team2 for j in range(num_tasks)]) <= team_max)
  sol = solver.Solve()

  print('Total cost = ', solver.Objective().Value())
  print()
  for i in range(num_workers):
    for j in range(num_tasks):
      if x[i, j].solution_value() > 0:
        print('Worker %d assigned to task %d.  Cost = %d' % (
              i,
              j,
              cost[i][j]))

  print()
  print("Time = ", solver.WallTime(), " milliseconds")
if __name__ == '__main__':
  main()