使用OR-Tools CP-SAT加速大规模指派问题求解

from ortools.sat.python import cp_model
import numpy as np

# number of workers and tasks
N = 40 # Increased to N=70 for demonstration of scalability

# cost table for each worker-task pairs
np.random.seed(0)
costs = np.random.rand(N,N)*100

# workers IDs    
workers_id = (np.random.rand(N)*4).astype(np.uint32)
id_2_idsrt_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D'}        
workers_id_str = [id_2_idsrt_dict[val] for val in workers_id]

idsrt_2_id_dict = {idstr: id for id, idstr in id_2_idsrt_dict.items()}

num_workers = len(costs)
num_tasks = len(costs[0])

# Scale costs to integers to use with CP-SAT
# Multiplying by 100 to keep two decimal places precision
scaling_factor = 100
scaled_costs = (costs * scaling_factor).astype(int)

# Determine bounds for scaled costs
min_scaled_cost_val = int(np.min(scaled_costs))
max_scaled_cost_val = int(np.max(scaled_costs))

# Max possible cost for a single worker (since each worker takes one task)
# This is also the upper bound for max_cost and lower bound for min_cost
max_possible_worker_cost = max_scaled_cost_val
min_possible_worker_cost = min_scaled_cost_val

# Create the CP-SAT model
model = cp_model.CpModel()

# Variables
# x[i, j] is a 0-1 variable, which will be 1 if worker i is assigned to task j.
x = {}
for i in range(num_workers):   
  for j in range(num_tasks):
     x[i, j] = model.NewBoolVar(f"x_{i}_{j}")

# tasks_ids[j] is an integer variable containing each task's assigned worker id.
# Worker IDs range from 0 to 3.
tasks_ids = []
for j in range(num_tasks):
   # The ID for a task will be the ID of the worker assigned to it.
   # We define its range from 0 to 3 (corresponding to 'A' to 'D').
   tasks_ids.append( model.NewIntVar(0, 3, f"task_id_{j}") )

   # Link task_id to the worker_id of the assigned worker
   # tasks_ids[j] == sum(workers_id[i] * x[i, j] for i in range(num_workers))
   model.Add(tasks_ids[j] == sum(workers_id[i] * x[i, j] for i in range(num_workers)))

# Constraint: Each worker is assigned to exactly one task.
for i in range(num_workers):
   model.Add(sum(x[i, j] for j in range(num_tasks)) == 1)

# Constraint: Each task is assigned to exactly one worker.
for j in range(num_tasks):
   model.Add(sum(x[i, j] for i in range(num_workers)) == 1)   

# Constraint: Task 1 can be assigned only with workers that have the id "A"     
model.Add(tasks_ids[1] == idsrt_2_id_dict["A"])  

# Constraint: Tasks 2,4,6 must assigned with workers of the same id
model.Add(tasks_ids[2] == tasks_ids[4]) 
model.Add(tasks_ids[2] == tasks_ids[6]) 

# Constraint: Tasks 10,11,12 must assigned with workers of the same id
model.Add(tasks_ids[10] == tasks_ids[11]) 
model.Add(tasks_ids[11] == tasks_ids[12]) 

# Constraint: Tasks 1,2,3 sum of ids <= 4
model.Add((tasks_ids[1] + tasks_ids[2] + tasks_ids[3]) <= 4) 

# Constraint: Tasks 4,5,6 sum of ids <= 4
model.Add((tasks_ids[4] + tasks_ids[5] + tasks_ids[6]) <= 4) 

# Constraint: Tasks 7,8,9 sum of ids <= 3
model.Add((tasks_ids[7] + tasks_ids[8] + tasks_ids[9]) <= 3) 

# Objective: minimize the difference of assignment higher cost worker and lower cost worker

# List of scaled costs for each worker's assignment
assignment_workers_scaled_costs_vars = []
for i in range(num_workers):            
   # Each worker's cost is the sum of scaled_costs[i][j] * x[i,j] for their assigned task.
   # The bounds for this variable are min_possible_worker_cost to max_possible_worker_cost.
   worker_cost_var = model.NewIntVar(min_possible_worker_cost, max_possible_worker_cost, f'worker_scaled_cost_{i}')
   model.Add(worker_cost_var == sum(scaled_costs[i][j] * x[i, j] for j in range(num_tasks)))
   assignment_workers_scaled_costs_vars.append(worker_cost_var)

# Additional variables for max and min costs
max_cost_var = model.NewIntVar(min_possible_worker_cost, max_possible_worker_cost, 'max_scaled_cost')        
min_cost_var = model.NewIntVar(min_possible_worker_cost, max_possible_worker_cost, 'min_scaled_cost')

# Constraints to update max and min costs using CP-SAT's dedicated functions
model.AddMaxEquality(max_cost_var, assignment_workers_scaled_costs_vars)
model.AddMinEquality(min_cost_var, assignment_workers_scaled_costs_vars)

# Minimize the difference between max and min costs
model.Minimize(max_cost_var - min_cost_var)

# Create a solver and solve the model.
solver = cp_model.CpSolver()

# Optional: Configure solver parameters for tuning
# solver.parameters.log_search_progress = True # Enable logging to see solver progress
# solver.parameters.num_workers = 8 # Use multiple cores for parallel search (if applicable)
# solver.parameters.max_time_in_seconds = 60.0 # Set a time limit

print(f"Solving with CP-SAT solver version: {solver.CpSolverVersion()}")
status = solver.Solve(model)

# Print solution.
if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
   # The objective value is scaled, so divide by scaling_factor to get original units
   objective_diff = solver.ObjectiveValue() / scaling_factor
   print(f"Difference (min_max_cost) = {objective_diff:.2f}\n")
   print(f"Max scaled cost: {solver.Value(max_cost_var) / scaling_factor:.2f}")
   print(f"Min scaled cost: {solver.Value(min_cost_var) / scaling_factor:.2f}")

   for i in range(num_workers):
      for j in range(num_tasks):
         if solver.Value(x[i, j]) > 0.5: # If x[i,j] is 1
            print(f"Worker {i} ({workers_id_str[i]}) assigned to task {j}." + 
                  f" Original Cost: {costs[i][j]:.2f}, Scaled Cost: {scaled_costs[i][j]}")
else:
   print("No solution found.")
   if status == cp_model.INFEASIBLE:
       print("The problem is infeasible.")
   elif status == cp_model.MODEL_INVALID:
       print("The model is invalid.")
   else:
       print(f"Solver status: {solver.StatusName(status)}")