preference-allocator.py

#!/usr/bin/env python3
import random
from ortools.sat.python import cp_model

    
def generatePreferences(n, m):
    '''
    n - teams
    m - customers
    Randomly Generates Preferences of customer for n teams.
    l[p][q] returns what team p ranked customer q 
    '''
    p = [i for i in range(m)]
    l = []
    for i in range(n):
        q = p[:]
        random.shuffle(q)
        l.append(q)
    return l

def main(n, m):
    preferences = generatePreferences(n, m)
    
    num_teams = n
    all_teams = range(n)
    num_customers = m
    all_customers = range(m)
    
    model = cp_model.CpModel()
    
    # Variables
    allocations = []
    for ti in all_teams:
        t = []
        for pi in all_customers:
            t.append(model.NewBoolVar('x[%i,%i]' % (ti, pi)))
        allocations.append(t)
    # Constraints
    
    # Each customer has atmost 2 teams
    [model.Add(sum(allocations[ti][pi] for ti in all_teams) ==2) for pi in all_customers]

    # Each team has one Customer
    [model.Add(sum(allocations[ti][pi] for pi in all_customers) == 1) for ti in all_teams]

    
    model.Minimize(sum([allocations[ti][pi] * preferences[ti][pi] for ti in all_teams for pi in all_customers]))
    
    solver = cp_model.CpSolver()
    status = solver.Solve(model)
    
    if status == cp_model.OPTIMAL:
        print('Mean Dissatisfaction :( = %s' % str(float(solver.ObjectiveValue())/num_teams))
        print()
        for ti in all_teams:
            for pi in all_customers:
                if solver.Value(allocations[ti][pi]) == 1:
                    print('Team ', ti, ' assigned Customer ', pi, '  Preference # = ',
                          preferences[ti][pi])

    print()

    print('Statistics')
    print('  - conflicts : %i' % solver.NumConflicts())
    print('  - branches  : %i' % solver.NumBranches())
    print('  - wall time : %f s' % solver.WallTime())


if __name__ == '__main__':
    print(generatePreferences(10, 5))
    main(10, 5)