CSE505/Assignment2/B/task.py
2024-03-09 23:16:18 +08:00

53 lines
1.6 KiB
Python

"""Solves a simple assignment problem."""
from ortools.sat.python import cp_model
import timeit
def main() -> None:
# Data
costs = [1, 19, 21, 8, 6, 24, 17, 23, 15, 22, 14, 2, 2, 10, 3, 17, 17, 19, 26, 9, 28, 16, 22, 19, 17, 15, 18, 15, 23, 27]
num_workers = 8
deadline = 90
num_tasks = len(costs)
# Model
model = cp_model.CpModel()
# Variables
x = {}
for worker in range(num_workers):
for task in range(num_tasks):
x[worker, task] = model.new_bool_var(f"x[{worker},{task}]")
# Each task is assigned to exactly one worker.
for task in range(num_tasks):
model.add_exactly_one(x[worker, task] for worker in range(num_workers))
# Objective
for worker in range(num_workers):
objective_terms = []
for task in range(num_tasks):
objective_terms.append(costs[task] * x[worker, task])
model.AddLinearConstraint(sum(objective_terms), 0, deadline)
# Solve
solver = cp_model.CpSolver()
status = solver.solve(model)
# Print solution.
if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
print(f"Total cost = {solver.objective_value}\n")
for worker in range(num_workers):
for task in range(num_tasks):
if solver.boolean_value(x[worker, task]):
print(
f"Worker {worker} assigned to task {task}."
+ f" Cost = {costs[task]}"
)
else:
print("No solution found.")
if __name__ == "__main__":
start = timeit.default_timer()
main()
stop = timeit.default_timer()
print('Time: ', stop - start)