Workplace scheduling with Python-MIP

schedule.py

# From https://www.python-mip.com/
from mip import *

m = Model(sense = MAXIMIZE)

num_slots = 10
num_workers = 5
num_tasks = 3

def slot_to_time(slot):
    hour = 10 + slot // 2
    minute = 30 if slot % 2 else 0
    return f'{hour:02d}:{minute:02d}'

# shift_slots[i][j] means shift i includes slot j

def make_shift(total, start, end):
    return start * [0] + (end - start) * [1] + (total - end) * [0]

short_shift_slots = [make_shift(num_slots, start, start + 2)
                     for start in range(0, num_slots - 1)]
shift_slots = (short_shift_slots +
               [make_shift(num_slots, start, start + 3)
                for start in range(0, num_slots - 2)])

lunch_slots = [make_shift(num_slots, start, start + 1)
               for start in range(3, 7)]

num_shifts = len(shift_slots)
num_short_shifts = len(short_shift_slots)
num_lunches = len(lunch_slots)

# assigment[worker][shift][task]

assigment = [[[m.add_var(var_type=BINARY)
               for task in range(num_tasks)]
              for shift in range(num_shifts)]
             for worker in range(num_workers)]

# lunch_assigment[worker][lunch]
lunch_assigment = [[m.add_var(var_type=BINARY)
                    for lunch in range(num_lunches)]
                   for worker in range(num_workers)]

# For every task, every slot has to be worked by some worker
for task in range(num_tasks):
    for slot in range(num_slots):
        m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
                for worker in range(num_workers)
                for shift in range(num_shifts)) == 1)

# The first task only allows short shifts
task = 0
for worker in range(num_workers):
    for shift in range(num_short_shifts, num_shifts):
        m += assigment[worker][shift][task] == 0

# Ensure no worker either works or has lunch more than one slot at one time.
for slot in range(num_slots):
    for worker in range(num_workers):
        m += (xsum(assigment[worker][shift][task] * shift_slots[shift][slot]
                   for shift in range(num_shifts)
                   for task in range(num_tasks)) +
              xsum(lunch_slots[lunch][slot] * lunch_assigment[worker][lunch]
                   for lunch in range(num_lunches))) <= 1

# Nobody is allowed to work more than two slots on task 0
task = 0
for worker in range(num_workers):
    m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
            for slot in range(num_slots)
            for shift in range(num_shifts)) <= 2)

# Nobody is allowed to work more than this many consecutive slots on one task
max_slots_on_one_task = 3
for worker in range(num_workers):
    for task in range(num_tasks):
        for slot_start in range(0, num_slots - max_slots_on_one_task):
            amount_worked = xsum(assigment[worker][shift][task] * shift_slots[shift][slot]
                                 for shift in range(num_shifts)
                                 for slot in range(slot_start, slot_start + max_slots_on_one_task + 1))
            m += amount_worked <= max_slots_on_one_task

# Nobody is allowed to work more this many consecutive slots
max_slots_on_any_task = 5
for worker in range(num_workers):
    for slot_start in range(0, num_slots - max_slots_on_any_task):
        amount_worked = xsum(assigment[worker][shift][task] * shift_slots[shift][slot]
                             for shift in range(num_shifts)
                             for task in range(num_tasks)
                             for slot in range(slot_start, slot_start + max_slots_on_any_task + 1))
        m += amount_worked <= max_slots_on_any_task

# Worker 0 works 7 slots
worker = 0
m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
        for task in range(num_tasks)
        for slot in range(num_slots)
        for shift in range(num_shifts)) == 7)
# Worker 1 works 7 slots
worker = 1
m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
        for task in range(num_tasks)
        for slot in range(num_slots)
        for shift in range(num_shifts)) == 7)
# Worker 2 works 5 slots
worker = 2
m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
        for task in range(num_tasks)
        for slot in range(num_slots)
        for shift in range(num_shifts)) == 5)
# Worker 3 works 5 slots
worker = 3
m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
        for task in range(num_tasks)
        for slot in range(num_slots)
        for shift in range(num_shifts)) == 5)
# Worker 4 works 6 slots
worker = 4
m += (xsum(shift_slots[shift][slot] * assigment[worker][shift][task]
        for task in range(num_tasks)
        for slot in range(num_slots)
        for shift in range(num_shifts)) == 6)

# Person 4 has 12:30 lunch
m += lunch_assigment[4][2] == 1

# Persons 2 and 3 can't lunch together
for slot in range(num_slots):
    m += xsum(lunch_slots[lunch][slot] * lunch_assigment[worker][lunch]
               for lunch in range(num_lunches)
               for worker in [2, 3]) <= 1

# Every worker has one lunch slot
for worker in range(num_workers):
    m += xsum(lunch_slots[lunch][slot] * lunch_assigment[worker][lunch]
               for lunch in range(num_lunches)
               for slot in range(num_slots)) == 1

m.objective = 1

print("Optimizing")

status = m.optimize()

print("status =", m.status)
# for task in range(num_tasks):
#     print('task ', task)
#     for shift in range(num_shifts):
#         for worker in range(num_workers):
#             print(assigment[worker][shift][task].x, ' ', end='')
#         print()
#     print('')

for slot in range(num_slots):
    print(slot_to_time(slot), end=' ')

    for task in range(num_tasks):
        assigned_worker = -1
        for worker in range(num_workers):
            is_worked_by = sum(assigment[worker][shift][task].x * shift_slots[shift][slot]
                               for shift in range(num_shifts))
            if is_worked_by:
                assigned_worker = worker
        assert(assigned_worker >= 0)
        print(assigned_worker, '', end='')

    print(' ', end='')

    for worker in range(num_workers):
        is_lunched_by = sum(lunch_assigment[worker][lunch].x * lunch_slots[lunch][slot]
                           for lunch in range(num_lunches))
        if is_lunched_by:
            print(worker, ' ', end='')

    print('')