优化OR-Tools解决大规模分配问题：CP-SAT的性能优势与浮点数缩放

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

# number of workers and tasks
N = 40 # 尝试更大的N值，例如70

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

# 为了CP-SAT更好地处理浮点数，我们可以将成本缩放到整数
# 例如，乘以1000，然后进行四舍五入
# CP-SAT内部也会进行类似操作，但明确转换有时能提供更好的控制或性能
SCALING_FACTOR = 1000
scaled_costs = (costs * SCALING_FACTOR).astype(int)

# 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]
# print(f"Worker IDs: {workers_id_str}") # 打印信息可以帮助调试

idsrt_2_id_dict = {}
for id_val, idstr in id_2_idsrt_dict.items():
   idsrt_2_id_dict[idstr] = id_val
# print(f"ID string to ID dict: {idsrt_2_id_dict}")

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

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

# Variables
# x[i, j] is an array of 0-1 variables, 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}") # 使用NewBoolVar创建布尔变量

# Variables
# tasks_ids[j] is a list of integers variables contains each task's assigned worker id.
# 这里需要NewIntVar，因为workers_id是整数
tasks_ids = []
for j in range(num_tasks):
   # task_id_j = sum([workers_id[i]*x[i, j] for i in range(num_workers)])
   # CP-SAT不能直接对NewBoolVar求和得到NewIntVar，需要使用model.NewIntVar和model.Add
   task_id_var = model.NewIntVar(0, max(workers_id), f"task_id_{j}")
   model.Add(task_id_var == sum([workers_id[i]*x[i, j] for i in range(num_workers)]))
   tasks_ids.append(task_id_var)

# 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 workers costs for an assignment
assignment_workers_costs_list = []
for i in range(num_workers):
   # 注意这里使用 scaled_costs
   worker_cost_var = model.NewIntVar(0, int(np.max(scaled_costs)), f"worker_cost_{i}")
   model.Add(worker_cost_var == sum([scaled_costs[i][j] * x[i, j] for j in range(num_tasks)]))
   assignment_workers_costs_list.append(worker_cost_var)

# Additional variables for max and min costs
# CP-SAT的Max和Min操作需要一个列表的变量
max_cost_var = model.NewIntVar(0, int(np.max(scaled_costs)), 'max_cost')
min_cost_var = model.NewIntVar(0, int(np.max(scaled_costs)), 'min_cost') # min_cost_limit在CP-SAT中通常设为0或最小可能值

# Constraints to update max and min costs
# 使用model.AddMaxEquality和model.AddMinEquality
model.AddMaxEquality(max_cost_var, assignment_workers_costs_list)
model.AddMinEquality(min_cost_var, assignment_workers_costs_list)

# 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()
# 可以设置一些参数来观察求解过程
# solver.parameters.log_search_progress = True
# solver.parameters.num_workers = 8 # 根据CPU核心数调整

print(f"Solving with CP-SAT {solver.parameters.solver_version}")
status = solver.Solve(model)

# Print solution.
if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
   # 目标值需要除以缩放因子
   objective_diff = solver.ObjectiveValue() / SCALING_FACTOR
   print(f"Difference (scaled back)= {objective_diff:.2f}\n")
   for i in range(num_workers):
      for j in range(num_tasks):
         if solver.Value(x[i, j]) > 0.5:
            # 成本也需要除以缩放因子
            original_cost = costs[i][j]
            print(f"Worker {i} ({workers_id_str[i]}) assigned to task {j}." + f" Cost: {original_cost:.2f}")
else:
   print("No solution found.")