diff --git a/examples/finished/categorical_data.py b/examples/finished/categorical_data.py index eb259015..07f122aa 100644 --- a/examples/finished/categorical_data.py +++ b/examples/finished/categorical_data.py @@ -1,15 +1,22 @@ +""" +This example shows how one can optimize a model with categorical data by converting it into integers. + +There are three employees (Alice, Bob, Charlie) and three shifts. Each shift is assigned an integer: + +Morning - 0 +Afternoon - 1 +Night - 2 + +The employees have availabilities (e.g. Alice can only work in the Morning and Afternoon), and different +salary demands. These constraints, and an additional one stipulating that every shift must be covered, +allows us to model a MIP with the objective of minimizing the money spent on salary. +""" + from pyscipopt import Model # Define categorical data -shifts = {"Morning": 0, "Afternoon": 1, "Night": 2} -employees = ["Alice", "Bob", "Charlie"] - -# Employees have different salary demands -cost = { - "Alice": [2,4,1], - "Bob": [3,2,7], - "Charlie": [3,3,3] -} +shift_to_int = {"Morning": 0, "Afternoon": 1, "Night": 2} +employees = ["Alice", "Bob", "Charlie"] # Employee availability availability = { @@ -19,9 +26,16 @@ } # Transform availability into integer values -availability_int = { - emp: [shifts[shift] for shift in available_shifts] - for emp, available_shifts in availability.items() +availability_int = {} +for emp, available_shifts in availability.items(): + availability_int[emp] = [shift_to_int[shift] for shift in available_shifts] + + +# Employees have different salary demands +cost = { + "Alice": [2,4,1], + "Bob": [3,2,7], + "Charlie": [3,3,3] } # Create the model @@ -30,22 +44,22 @@ # x[e, s] = 1 if employee e is assigned to shift s x = {} for e in employees: - for s in shifts.values(): + for s in shift_to_int.values(): x[e, s] = model.addVar(vtype="B", name=f"x({e},{s})") # Each shift must be assigned to exactly one employee -for s in shifts.values(): +for s in shift_to_int.values(): model.addCons(sum(x[e, s] for e in employees) == 1) # Employees can only work shifts they are available for for e in employees: - for s in shifts.values(): + for s in shift_to_int.values(): if s not in availability_int[e]: model.addCons(x[e, s] == 0) # Minimize shift assignment cost model.setObjective( - sum(cost[e][s]*x[e, s] for e in employees for s in shifts.values()), "minimize" + sum(cost[e][s]*x[e, s] for e in employees for s in shift_to_int.values()), "minimize" ) # Solve the problem @@ -54,6 +68,6 @@ # Display the results print("\nOptimal Shift Assignment:") for e in employees: - for s, s_id in shifts.items(): + for s, s_id in shift_to_int.items(): if model.getVal(x[e, s_id]) > 0.5: print("%s is assigned to %s" % (e, s))