CSE505/Assignment2/B/task.py
2024-03-09 16:25:30 -05:00

51 lines
1.4 KiB
Python

"""Solves a simple assignment problem."""
from ortools.sat.python import cp_model
def main() -> None:
# Data
costs = [10,5,8,10,5,2,1]
num_workers = 2
deadline = 20
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__":
main()