Constraint on uniqueness for Assignement Optimization Problem using OR-Tools CP-SAT Solver

Question

The problem I am trying to solve is pretty similar to the one described in the OR-tools documentation: https://developers.google.com/optimization/assignment/assignment_example#cp_sat_solution. I have 15 workers that I need to assign to 5 tasks. The objective function is to minimize the cost. However, the selected workers need to belong to exactly 3 states. I am having an issue expressing this constraint. Code below

-- Edited with solution

start_time = time.perf_counter()
costs = [500, 75, 80, 25, 95, 100, 10, 30, 400, 50, 80, 110, 150, 20, 60]
# states_label = ["AZ","CT","FL","CA","NY","CA","NY","CT","AZ","CA","CT","CT","AZ","FL","NY"]
states = [0,2,3,1,4,1,4,2,0,1,2,2,0,3,4]

num_workers = len(costs)
num_tasks = 5
num_states = 5

# Model
model = cp_model.CpModel()

# Variables
x = {}
for worker in range(num_workers):
    for task in range(num_tasks):
        x[worker, task] = model.NewBoolVar(f'x[{worker},{task}]')

# Create 
states_is_selected = []
for state in range(num_states):
    state_is_selected = model.NewBoolVar('')
    # workers_from_that_states are all x that corresponds to this state
    workers_from_that_state = []
    for worker in range(num_workers):
        for task in range(num_tasks):
            if (state == states[worker]):
                workers_from_that_state.append( x[worker,task])
    # print(f"state {state}")
    # print(workers_from_that_state)
    literals = [c for c in workers_from_that_state]
    model.AddBoolOr(literals).OnlyEnforceIf(state_is_selected)
    for literal in literals:
        model.AddImplication(literal, state_is_selected)
    
    states_is_selected.append(state_is_selected)


# Constraints

# Each worker is assigned to at most 1 task.
for worker in range(num_workers):
    model.AddAtMostOne(x[worker,task] for task in range(num_tasks))

# Each task is assigned to exactly one worker.
for task in range(num_tasks):
    model.AddExactlyOne(x[worker,task] for worker in range(num_workers))

# The workers belong to exactly 3 states
model.Add(sum(states_is_selected) == 3)


# Objective
objective_terms = []
for worker in range(num_workers):
    for task in range(num_tasks):
        objective_terms.append(costs[worker] * x[worker,task])

model.Minimize(sum(objective_terms))

# Solve
solver = cp_model.CpSolver()
status = solver.Solve(model)


print(f'Status = {solver.StatusName(status)}')
if status == cp_model.INFEASIBLE:
    print('SufficientAssumptionsForInfeasibility = 'f'{solver.SufficientAssumptionsForInfeasibility()}')

# Print solution.
if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
    
    print(f"Solution found")
    
    print(f'Total cost = {solver.ObjectiveValue()}')
    print()
    for worker in range(num_workers):
        for task in range(num_tasks):
            if solver.BooleanValue(x[worker,task]):
                print(f"Worker {worker}, Cost: {costs[worker]}, State: {states[worker]}")   

    for state in range(num_states):
        print(f"State {state}: ", solver.BooleanValue(states_is_selected[state]))

    # Number of States for assigned workers
    states_selected = [states[worker] for worker in range(num_workers) for task in range(num_tasks) if solver.BooleanValue(x[worker,task])]
    print(f"The workers selected are coming from {len(list(set(states_selected)))} states:", set(states_selected))


else:  
    print(f"No solution found")

    
end_time = time.perf_counter()
print("Program executed in {:,.2f} s".format(end_time - start_time))

Answer 1

You need to have 1 Boolean variable per state.

Then implement

state_is_selected <=> bool_or(workers_from_that_state)

This is a simple generalization of: https://github.com/google/or-tools/blob/stable/ortools/sat/docs/boolean_logic.md#product-of-two-boolean-variables

then sum(states) == 3