pythonpandasalgorithmhungarian-algorithm

Hungarian Algorithm Constraint using Python

I have dataframe of job and employee, with the duration each employee can finish each job. I want to use Hungarian algorithm to assign each job to 1 employee, and each employee can only assigned by 1 job.

Here is the data:

    J1   J2   J3
E1  3    5    2
E2  7    3    9
E3  4    2    7
E4  2    7    5
E5  6    8    8

The expected result was:

    J1   J2   J3
E1  0    0    1
E2  0    0    0
E3  0    1    0
E4  1    0    0
E5  0    0    0

And then print out this:

Job   Employee
 1       4
 2       3
 3       1

Can anyone help me on this matter? Thanks in advance!

Solution

  • Assuming df the input dataframe, you can use scipy.optimize.linear_sum_assignment:

    from scipy.optimize import linear_sum_assignment
x,y = linear_sum_assignment(df)
out = pd.DataFrame({'Job': df.columns[y], 'Employee': df.index[x]})

    output:

      Job Employee
0  J3       E1
1  J2       E3
2  J1       E4