on-call generator

hungarian.py

import json
from scipy.optimize import linear_sum_assignment
import numpy as np

maximize = False # set to true if you want to maximize cost vs minimize (reverse order of preferences)

holidays = [
    'Memorial',
    'Independence Day',
    'Labor',
    'Thanksgiving',
    'Christmas',
    'New Years',
]

preferences = {
    'Person1': [
        'Memorial',
        'Labor',
        'Thanksgiving',
        'Independence Day',
        'Christmas',
        'New Years',
    ],
    'Person2': [
       'Thanksgiving',
        'Memorial',
        'Independence Day',
        'Labor',
        'Christmas',
        'New Years',
    ],
    'Person3': [
        'Thanksgiving',
        'Labor',
        'Christmas',
        'Memorial',
        'New Years',
        'Independence Day',
    ],
    'Person4': [
        'Thanksgiving',
        'Christmas',
        'Labor',
        'Memorial',
        'Independence Day',
        'New Years',
    ],

    # pad out with people who volunteered for extra days - number of people in list needs to match number of holidays available
    'Person1-2': [
        'Memorial',
        'Labor',
        'Independence Day',
        'Christmas',
        'Thanksgiving',
        'New Years',
    ],
    'Person2-2': [
        'Thanksgiving',
        'Memorial',
        'Independence Day',
        'Labor',
        'Christmas',
        'New Years',
    ],
}

if len(preferences.keys()) != len(list(preferences.values())[0]):
    quit("Be sure to pad out free days with extra members")

matrix = []

for person in list(preferences.keys()):
    person_preference_list = list(preferences[person])
    if maximize == True:
        person_preference_list.reverse()

    person_preference = list(map(lambda x : preferences[person].index(x) + 1, holidays))

    matrix.append(person_preference)

matrix = np.array(matrix)

row_indices, col_indices = linear_sum_assignment(matrix, maximize=maximize)

# print results
row_names = list(preferences.keys())
col_names = holidays

output = [((row_names[r], col_names[c]), matrix[r, c]) for r, c in zip(row_indices, col_indices)]
print(output)