EVA_duty_arrange_tool/solve.py

143 lines
5.2 KiB
Python

from ortools.sat.python import cp_model
from ortools.linear_solver import pywraplp
from utils import read_excel, save_to_excel
def solve_program(preference_mat:list,
want_num_array:list,
is_new_array:list,
is_tech_array:list,
is_hr_array:list,
num_min=None,
num_max=None,
num_tech_min=None,
num_tech_max=None,
num_old_min=None,
num_old_max=None,
num_hr_min=None,
num_hr_max=None,
num_new_min=None,
num_new_max=None
):
is_old_array = [not is_new for is_new in is_new_array]
is_not_tech_array = [not is_tech for is_tech in is_tech_array]
# 这是一个组合优化问题,我们使用 SCIP 求解器来解决这个问题
solver = pywraplp.Solver.CreateSolver("SCIP")
if not solver:
return
# 定义变量的规模,一共有 N*M 个变量需要求解器求解
N = len(preference_mat) # 学生人数
M = len(preference_mat[0]) # 班次数
avg_num = sum(want_num_array) / M # 每班次的平均人数
print(f"平均人数:{avg_num}")
variables = []
aux_vars = [] # 辅助变量
infinity = solver.infinity()
for i in range(N):
row_vars = []
for j in range(M):
var = solver.IntVar(0.0, 1.0, f"choice_{i}_{j}")
row_vars.append(var)
variables.append(row_vars)
for j in range(M):
aux_var = solver.NumVar(0.0, infinity, f"aux_{j}")
aux_vars.append(aux_var)
print("Number of variables =", solver.NumVariables())
# 添加约束:每个同学意愿一定要满足
for i in range(N):
for j in range(M):
solver.Add(variables[i][j] <= preference_mat[i][j])
# 辅助变量添加约束:
for j in range(M):
actual_num = sum(variables[i][j] for i in range(N)) # 每班次的实际人数
solver.Add(aux_vars[j] >= actual_num-avg_num)
solver.Add(aux_vars[j] >= avg_num-actual_num)
# 添加约束:满足每位同学的意愿班次
for i in range(N):
solver.Add(sum(variables[i]) == want_num_array[i])
# 添加约束:每个班次至少有?位同学
if num_min is not None:
for j in range(M):
solver.Add(sum(variables[i][j] for i in range(N)) >= num_min)
# 添加约束:每个班次至多有?位同学
if num_max is not None:
for j in range(M):
solver.Add(sum(variables[i][j] for i in range(N)) <= num_max)
# 添加约束:每班次最少包含?个电脑或电器的老人
if num_tech_min is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_old_array[i]*is_tech_array[i] for i in range(N)) >= num_tech_min)
# 添加约束:每班次最多包含?个电脑或电器的老人
if num_tech_max is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_old_array[i]*is_tech_array[i] for i in range(N)) <= num_tech_max)
# 添加约束:每班次至少包含?个老人
if num_old_min is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_old_array[i] for i in range(N)) >= num_old_min)
# 添加约束:每班次至多包含?个老人
if num_old_max is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_old_array[i] for i in range(N)) <= num_old_max)
# 添加约束:每班次至少包含?个人资部小朋友
if num_hr_min is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_hr_array[i] for i in range(N)) >= num_hr_min)
# 添加约束:每班次至多包含?个人资部小朋友
if num_hr_max is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_hr_array[i] for i in range(N)) <= num_hr_max)
# 添加约束:每班次至少包含?个小朋友
if num_new_min is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_new_array[i] for i in range(N)) >= num_new_min)
# 添加约束:每班次至多包含?个小朋友
if num_new_max is not None:
for j in range(M):
solver.Add(sum(variables[i][j]*is_new_array[i] for i in range(N)) <= num_new_max)
# 优化目标:每班次人数尽可能平均
solver.Minimize(sum(aux_vars)) # Maximize the sum of these variables.
# 求解优化问题
status = solver.Solve()
# 输出结果
variables_return = []
aux_vars_return = []
if status == pywraplp.Solver.OPTIMAL:
print("Optimal solution found:")
for i in range(N):
row_solution = [variables[i][j].solution_value() for j in range(M)]
variables_return.append(row_solution)
for j in range(M):
aux_vars_return.append(aux_vars[j].solution_value())
# Print the optimized value of the objective function.
print(f"Optimized objective value: {solver.Objective().Value()}")
print(aux_vars_return)
return variables_return
else:
print("No solution found.")
return None