CSE505/Assignment2/B/task.py

"""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)