#!/usr/bin/env python3
# Copyright (c) 2016 Anish Athalye.
# Released under the MIT License.
import csv
import sys
import itertools
import scipy.optimize
import numpy as np
# INPUT: CSV, in the format:
# "Name", "Group 1", "Group 2", "Group 3", ...
# ---, desired number of people in group 1, desired number of people in group 2, ...
# person 1 name, ranking of group 1, ranking of group 2, ...
# person 2 name, ranking of group 1, ranking of group 2, ...
# ...
#
# pass in the path to the CSV file as argv[1]
def main():
with open(sys.argv[1]) as csvfile:
rows = list(csv.reader(csvfile))
for row in rows:
assert len(row) == len(rows[0])
headers = rows[0]
capacities = rows[1]
people = rows[2:]
assert sum(map(int, capacities[1:])) == len(people)
assignments, cost = optimal_assignment(headers, capacities, people)
snd = lambda i: i[1]
bygroup = sorted(assignments, key=snd)
groups = []
for k, g in itertools.groupby(bygroup, key=snd):
print("\n=== %s ===" % k)
print('\n'.join('%s (cost %.1f)' % (i[0], i[2]) for i in g))
print("\n(cost: %.1f)" % cost)
cost_buckets = {}
for i in assignments:
cost = i[2]
cost_buckets[cost] = cost_buckets.get(cost, 0) + 1
for bucket in sorted(cost_buckets.keys()):
print("(%d people got their #%d choice)" % (cost_buckets[bucket], bucket))
def optimal_assignment(headers, capacities, people):
C = np.zeros((len(people), len(people)), dtype=np.int)
for i_person in range(len(people)):
j = 0
for i_group in range(len(headers[1:])):
capacity = int(capacities[i_group + 1])
C[i_person, j:(j+capacity)] = int(people[i_person][i_group + 1])
j += capacity
group_name_map = {}
j = 0
for i_group in range(len(headers[1:])):
capacity = int(capacities[i_group + 1])
for k in range(j, j + capacity):
group_name_map[k] = headers[i_group + 1]
j += capacity
row_ind, col_ind = scipy.optimize.linear_sum_assignment(C)
cost = C[row_ind, col_ind].sum()
assert len(row_ind) == len(people)
assignments = []
for r, c in zip(row_ind, col_ind):
assignments.append((people[r][0], group_name_map[c], C[r][c]))
return assignments, cost
if __name__ == '__main__':
if not hasattr(scipy.optimize, 'linear_sum_assignment'):
print("You need at least SciPy 0.17.0 (which is when linear_sum_assignment was added)")
exit(1)
main()