Module munkres
[hide private]
[frames] | no frames]

Source Code for Module munkres

  1  #!/usr/bin/env python 
  2  # -*- coding: iso-8859-1 -*- 
  3   
  4  # Documentation is intended to be processed by Epydoc. 
  5   
  6  """ 
  7  Introduction 
  8  ============ 
  9   
 10  The Munkres module provides an implementation of the Munkres algorithm 
 11  (also called the Hungarian algorithm or the Kuhn-Munkres algorithm), 
 12  useful for solving the Assignment Problem. 
 13   
 14  Assignment Problem 
 15  ================== 
 16   
 17  Let *C* be an *n*\ x\ *n* matrix representing the costs of each of *n* workers 
 18  to perform any of *n* jobs. The assignment problem is to assign jobs to 
 19  workers in a way that minimizes the total cost. Since each worker can perform 
 20  only one job and each job can be assigned to only one worker the assignments 
 21  represent an independent set of the matrix *C*. 
 22   
 23  One way to generate the optimal set is to create all permutations of 
 24  the indexes necessary to traverse the matrix so that no row and column 
 25  are used more than once. For instance, given this matrix (expressed in 
 26  Python):: 
 27   
 28      matrix = [[5, 9, 1], 
 29                [10, 3, 2], 
 30                [8, 7, 4]] 
 31   
 32  You could use this code to generate the traversal indexes:: 
 33   
 34      def permute(a, results): 
 35          if len(a) == 1: 
 36              results.insert(len(results), a) 
 37   
 38          else: 
 39              for i in range(0, len(a)): 
 40                  element = a[i] 
 41                  a_copy = [a[j] for j in range(0, len(a)) if j != i] 
 42                  subresults = [] 
 43                  permute(a_copy, subresults) 
 44                  for subresult in subresults: 
 45                      result = [element] + subresult 
 46                      results.insert(len(results), result) 
 47   
 48      results = [] 
 49      permute(range(len(matrix)), results) # [0, 1, 2] for a 3x3 matrix 
 50   
 51  After the call to permute(), the results matrix would look like this:: 
 52   
 53      [[0, 1, 2], 
 54       [0, 2, 1], 
 55       [1, 0, 2], 
 56       [1, 2, 0], 
 57       [2, 0, 1], 
 58       [2, 1, 0]] 
 59   
 60  You could then use that index matrix to loop over the original cost matrix 
 61  and calculate the smallest cost of the combinations:: 
 62   
 63      n = len(matrix) 
 64      minval = sys.maxint 
 65      for row in range(n): 
 66          cost = 0 
 67          for col in range(n): 
 68              cost += matrix[row][col] 
 69          minval = min(cost, minval) 
 70   
 71      print minval 
 72   
 73  While this approach works fine for small matrices, it does not scale. It 
 74  executes in O(*n*!) time: Calculating the permutations for an *n*\ x\ *n* 
 75  matrix requires *n*! operations. For a 12x12 matrix, that's 479,001,600 
 76  traversals. Even if you could manage to perform each traversal in just one 
 77  millisecond, it would still take more than 133 hours to perform the entire 
 78  traversal. A 20x20 matrix would take 2,432,902,008,176,640,000 operations. At 
 79  an optimistic millisecond per operation, that's more than 77 million years. 
 80   
 81  The Munkres algorithm runs in O(*n*\ ^3) time, rather than O(*n*!). This 
 82  package provides an implementation of that algorithm. 
 83   
 84  This version is based on 
 85  http://www.public.iastate.edu/~ddoty/HungarianAlgorithm.html. 
 86   
 87  This version was written for Python by Brian Clapper from the (Ada) algorithm 
 88  at the above web site. (The ``Algorithm::Munkres`` Perl version, in CPAN, was 
 89  clearly adapted from the same web site.) 
 90   
 91  Usage 
 92  ===== 
 93   
 94  Construct a Munkres object:: 
 95   
 96      from munkres import Munkres 
 97   
 98      m = Munkres() 
 99   
100  Then use it to compute the lowest cost assignment from a cost matrix. Here's 
101  a sample program:: 
102   
103      from munkres import Munkres, print_matrix 
104   
105      matrix = [[5, 9, 1], 
106                [10, 3, 2], 
107                [8, 7, 4]] 
108      m = Munkres() 
109      indexes = m.compute(matrix) 
110      print_matrix(matrix, msg='Lowest cost through this matrix:') 
111      total = 0 
112      for row, column in indexes: 
113          value = matrix[row][column] 
114          total += value 
115          print '(%d, %d) -> %d' % (row, column, value) 
116      print 'total cost: %d' % total 
117   
118  Running that program produces:: 
119   
120      Lowest cost through this matrix: 
121      [5, 9, 1] 
122      [10, 3, 2] 
123      [8, 7, 4] 
124      (0, 0) -> 5 
125      (1, 1) -> 3 
126      (2, 2) -> 4 
127      total cost=12 
128   
129  The instantiated Munkres object can be used multiple times on different 
130  matrices. 
131   
132  Non-square Cost Matrices 
133  ======================== 
134   
135  The Munkres algorithm assumes that the cost matrix is square. However, it's 
136  possible to use a rectangular matrix if you first pad it with 0 values to make 
137  it square. This module automatically pads rectangular cost matrices to make 
138  them square. 
139   
140  Notes: 
141   
142  - The module operates on a *copy* of the caller's matrix, so any padding will 
143    not be seen by the caller. 
144  - The cost matrix must be rectangular or square. An irregular matrix will 
145    *not* work. 
146   
147  Calculating Profit, Rather than Cost 
148  ==================================== 
149   
150  The cost matrix is just that: A cost matrix. The Munkres algorithm finds 
151  the combination of elements (one from each row and column) that results in 
152  the smallest cost. It's also possible to use the algorithm to maximize 
153  profit. To do that, however, you have to convert your profit matrix to a 
154  cost matrix. The simplest way to do that is to subtract all elements from a 
155  large value. For example:: 
156   
157      from munkres import Munkres, print_matrix 
158   
159      matrix = [[5, 9, 1], 
160                [10, 3, 2], 
161                [8, 7, 4]] 
162      cost_matrix = [] 
163      for row in matrix: 
164          cost_row = [] 
165          for col in row: 
166              cost_row += [sys.maxint - col] 
167          cost_matrix += [cost_row] 
168   
169      m = Munkres() 
170      indexes = m.compute(cost_matrix) 
171      print_matrix(matrix, msg='Highest profit through this matrix:') 
172      total = 0 
173      for row, column in indexes: 
174          value = matrix[row][column] 
175          total += value 
176          print '(%d, %d) -> %d' % (row, column, value) 
177   
178      print 'total profit=%d' % total 
179   
180  Running that program produces:: 
181   
182      Highest profit through this matrix: 
183      [5, 9, 1] 
184      [10, 3, 2] 
185      [8, 7, 4] 
186      (0, 1) -> 9 
187      (1, 0) -> 10 
188      (2, 2) -> 4 
189      total profit=23 
190   
191  The ``munkres`` module provides a convenience method for creating a cost 
192  matrix from a profit matrix. Since it doesn't know whether the matrix contains 
193  floating point numbers, decimals, or integers, you have to provide the 
194  conversion function; but the convenience method takes care of the actual 
195  creation of the cost matrix:: 
196   
197      import munkres 
198   
199      cost_matrix = munkres.make_cost_matrix(matrix, 
200                                             lambda cost: sys.maxint - cost) 
201   
202  So, the above profit-calculation program can be recast as:: 
203   
204      from munkres import Munkres, print_matrix, make_cost_matrix 
205   
206      matrix = [[5, 9, 1], 
207                [10, 3, 2], 
208                [8, 7, 4]] 
209      cost_matrix = make_cost_matrix(matrix, lambda cost: sys.maxint - cost) 
210      m = Munkres() 
211      indexes = m.compute(cost_matrix) 
212      print_matrix(matrix, msg='Lowest cost through this matrix:') 
213      total = 0 
214      for row, column in indexes: 
215          value = matrix[row][column] 
216          total += value 
217          print '(%d, %d) -> %d' % (row, column, value) 
218      print 'total profit=%d' % total 
219   
220  References 
221  ========== 
222   
223  1. http://www.public.iastate.edu/~ddoty/HungarianAlgorithm.html 
224   
225  2. Harold W. Kuhn. The Hungarian Method for the assignment problem. 
226     *Naval Research Logistics Quarterly*, 2:83-97, 1955. 
227   
228  3. Harold W. Kuhn. Variants of the Hungarian method for assignment 
229     problems. *Naval Research Logistics Quarterly*, 3: 253-258, 1956. 
230   
231  4. Munkres, J. Algorithms for the Assignment and Transportation Problems. 
232     *Journal of the Society of Industrial and Applied Mathematics*, 
233     5(1):32-38, March, 1957. 
234   
235  5. http://en.wikipedia.org/wiki/Hungarian_algorithm 
236   
237  Copyright and License 
238  ===================== 
239   
240  This software is released under a BSD license, adapted from 
241  <http://opensource.org/licenses/bsd-license.php> 
242   
243  Copyright (c) 2008 Brian M. Clapper 
244  All rights reserved. 
245   
246  Redistribution and use in source and binary forms, with or without 
247  modification, are permitted provided that the following conditions are met: 
248   
249  * Redistributions of source code must retain the above copyright notice, 
250    this list of conditions and the following disclaimer. 
251   
252  * Redistributions in binary form must reproduce the above copyright notice, 
253    this list of conditions and the following disclaimer in the documentation 
254    and/or other materials provided with the distribution. 
255   
256  * Neither the name "clapper.org" nor the names of its contributors may be 
257    used to endorse or promote products derived from this software without 
258    specific prior written permission. 
259   
260  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
261  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
262  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
263  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 
264  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
265  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
266  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
267  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
268  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
269  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
270  POSSIBILITY OF SUCH DAMAGE. 
271  """ 
272   
273  __docformat__ = 'restructuredtext' 
274   
275  # --------------------------------------------------------------------------- 
276  # Imports 
277  # --------------------------------------------------------------------------- 
278   
279  import sys 
280   
281  # --------------------------------------------------------------------------- 
282  # Exports 
283  # --------------------------------------------------------------------------- 
284   
285  __all__     = ['Munkres', 'make_cost_matrix'] 
286   
287  # --------------------------------------------------------------------------- 
288  # Globals 
289  # --------------------------------------------------------------------------- 
290   
291  # Info about the module 
292  __version__   = "1.0.5.4" 
293  __author__    = "Brian Clapper, bmc@clapper.org" 
294  __url__       = "http://bmc.github.com/munkres/" 
295  __copyright__ = "(c) 2008 Brian M. Clapper" 
296  __license__   = "BSD-style license" 
297   
298  # --------------------------------------------------------------------------- 
299  # Classes 
300  # --------------------------------------------------------------------------- 
301   
302 -class Munkres:
303 """ 304 Calculate the Munkres solution to the classical assignment problem. 305 See the module documentation for usage. 306 """ 307
308 - def __init__(self):
309 """Create a new instance""" 310 self.C = None 311 self.row_covered = [] 312 self.col_covered = [] 313 self.n = 0 314 self.Z0_r = 0 315 self.Z0_c = 0 316 self.marked = None 317 self.path = None
318
319 - def make_cost_matrix(profit_matrix, inversion_function):
320 """ 321 **DEPRECATED** 322 323 Please use the module function ``make_cost_matrix()``. 324 """ 325 import munkres 326 return munkres.make_cost_matrix(profit_matrix, inversion_function)
327 328 make_cost_matrix = staticmethod(make_cost_matrix) 329
330 - def pad_matrix(self, matrix, pad_value=0):
331 """ 332 Pad a possibly non-square matrix to make it square. 333 334 :Parameters: 335 matrix : list of lists 336 matrix to pad 337 338 pad_value : int 339 value to use to pad the matrix 340 341 :rtype: list of lists 342 :return: a new, possibly padded, matrix 343 """ 344 max_columns = 0 345 total_rows = len(matrix) 346 347 for row in matrix: 348 max_columns = max(max_columns, len(row)) 349 350 total_rows = max(max_columns, total_rows) 351 352 new_matrix = [] 353 for row in matrix: 354 row_len = len(row) 355 new_row = row[:] 356 if total_rows > row_len: 357 # Row too short. Pad it. 358 new_row += [0] * (total_rows - row_len) 359 new_matrix += [new_row] 360 361 while len(new_matrix) < total_rows: 362 new_matrix += [[0] * total_rows] 363 364 return new_matrix
365
366 - def compute(self, cost_matrix):
367 """ 368 Compute the indexes for the lowest-cost pairings between rows and 369 columns in the database. Returns a list of (row, column) tuples 370 that can be used to traverse the matrix. 371 372 :Parameters: 373 cost_matrix : list of lists 374 The cost matrix. If this cost matrix is not square, it 375 will be padded with zeros, via a call to ``pad_matrix()``. 376 (This method does *not* modify the caller's matrix. It 377 operates on a copy of the matrix.) 378 379 **WARNING**: This code handles square and rectangular 380 matrices. It does *not* handle irregular matrices. 381 382 :rtype: list 383 :return: A list of ``(row, column)`` tuples that describe the lowest 384 cost path through the matrix 385 386 """ 387 self.C = self.pad_matrix(cost_matrix) 388 self.n = len(self.C) 389 self.original_length = len(cost_matrix) 390 self.original_width = len(cost_matrix[0]) 391 self.row_covered = [False for i in range(self.n)] 392 self.col_covered = [False for i in range(self.n)] 393 self.Z0_r = 0 394 self.Z0_c = 0 395 self.path = self.__make_matrix(self.n * 2, 0) 396 self.marked = self.__make_matrix(self.n, 0) 397 398 done = False 399 step = 1 400 401 steps = { 1 : self.__step1, 402 2 : self.__step2, 403 3 : self.__step3, 404 4 : self.__step4, 405 5 : self.__step5, 406 6 : self.__step6 } 407 408 while not done: 409 try: 410 func = steps[step] 411 step = func() 412 except KeyError: 413 done = True 414 415 # Look for the starred columns 416 results = [] 417 for i in range(self.original_length): 418 for j in range(self.original_width): 419 if self.marked[i][j] == 1: 420 results += [(i, j)] 421 422 return results
423
424 - def __copy_matrix(self, matrix):
425 """Return an exact copy of the supplied matrix""" 426 return copy.deepcopy(matrix)
427
428 - def __make_matrix(self, n, val):
429 """Create an *n*x*n* matrix, populating it with the specific value.""" 430 matrix = [] 431 for i in range(n): 432 matrix += [[val for j in range(n)]] 433 return matrix
434
435 - def __step1(self):
436 """ 437 For each row of the matrix, find the smallest element and 438 subtract it from every element in its row. Go to Step 2. 439 """ 440 C = self.C 441 n = self.n 442 for i in range(n): 443 minval = min(self.C[i]) 444 # Find the minimum value for this row and subtract that minimum 445 # from every element in the row. 446 for j in range(n): 447 self.C[i][j] -= minval 448 449 return 2
450
451 - def __step2(self):
452 """ 453 Find a zero (Z) in the resulting matrix. If there is no starred 454 zero in its row or column, star Z. Repeat for each element in the 455 matrix. Go to Step 3. 456 """ 457 n = self.n 458 for i in range(n): 459 for j in range(n): 460 if (self.C[i][j] == 0) and \ 461 (not self.col_covered[j]) and \ 462 (not self.row_covered[i]): 463 self.marked[i][j] = 1 464 self.col_covered[j] = True 465 self.row_covered[i] = True 466 467 self.__clear_covers() 468 return 3
469
470 - def __step3(self):
471 """ 472 Cover each column containing a starred zero. If K columns are 473 covered, the starred zeros describe a complete set of unique 474 assignments. In this case, Go to DONE, otherwise, Go to Step 4. 475 """ 476 n = self.n 477 count = 0 478 for i in range(n): 479 for j in range(n): 480 if self.marked[i][j] == 1: 481 self.col_covered[j] = True 482 count += 1 483 484 if count >= n: 485 step = 7 # done 486 else: 487 step = 4 488 489 return step
490
491 - def __step4(self):
492 """ 493 Find a noncovered zero and prime it. If there is no starred zero 494 in the row containing this primed zero, Go to Step 5. Otherwise, 495 cover this row and uncover the column containing the starred 496 zero. Continue in this manner until there are no uncovered zeros 497 left. Save the smallest uncovered value and Go to Step 6. 498 """ 499 step = 0 500 done = False 501 row = -1 502 col = -1 503 star_col = -1 504 while not done: 505 (row, col) = self.__find_a_zero() 506 if row < 0: 507 done = True 508 step = 6 509 else: 510 self.marked[row][col] = 2 511 star_col = self.__find_star_in_row(row) 512 if star_col >= 0: 513 col = star_col 514 self.row_covered[row] = True 515 self.col_covered[col] = False 516 else: 517 done = True 518 self.Z0_r = row 519 self.Z0_c = col 520 step = 5 521 522 return step
523
524 - def __step5(self):
525 """ 526 Construct a series of alternating primed and starred zeros as 527 follows. Let Z0 represent the uncovered primed zero found in Step 4. 528 Let Z1 denote the starred zero in the column of Z0 (if any). 529 Let Z2 denote the primed zero in the row of Z1 (there will always 530 be one). Continue until the series terminates at a primed zero 531 that has no starred zero in its column. Unstar each starred zero 532 of the series, star each primed zero of the series, erase all 533 primes and uncover every line in the matrix. Return to Step 3 534 """ 535 count = 0 536 path = self.path 537 path[count][0] = self.Z0_r 538 path[count][1] = self.Z0_c 539 done = False 540 while not done: 541 row = self.__find_star_in_col(path[count][1]) 542 if row >= 0: 543 count += 1 544 path[count][0] = row 545 path[count][1] = path[count-1][1] 546 else: 547 done = True 548 549 if not done: 550 col = self.__find_prime_in_row(path[count][0]) 551 count += 1 552 path[count][0] = path[count-1][0] 553 path[count][1] = col 554 555 self.__convert_path(path, count) 556 self.__clear_covers() 557 self.__erase_primes() 558 return 3
559
560 - def __step6(self):
561 """ 562 Add the value found in Step 4 to every element of each covered 563 row, and subtract it from every element of each uncovered column. 564 Return to Step 4 without altering any stars, primes, or covered 565 lines. 566 """ 567 minval = self.__find_smallest() 568 for i in range(self.n): 569 for j in range(self.n): 570 if self.row_covered[i]: 571 self.C[i][j] += minval 572 if not self.col_covered[j]: 573 self.C[i][j] -= minval 574 return 4
575
576 - def __find_smallest(self):
577 """Find the smallest uncovered value in the matrix.""" 578 minval = sys.maxint 579 for i in range(self.n): 580 for j in range(self.n): 581 if (not self.row_covered[i]) and (not self.col_covered[j]): 582 if minval > self.C[i][j]: 583 minval = self.C[i][j] 584 return minval
585
586 - def __find_a_zero(self):
587 """Find the first uncovered element with value 0""" 588 row = -1 589 col = -1 590 i = 0 591 n = self.n 592 done = False 593 594 while not done: 595 j = 0 596 while True: 597 if (self.C[i][j] == 0) and \ 598 (not self.row_covered[i]) and \ 599 (not self.col_covered[j]): 600 row = i 601 col = j 602 done = True 603 j += 1 604 if j >= n: 605 break 606 i += 1 607 if i >= n: 608 done = True 609 610 return (row, col)
611
612 - def __find_star_in_row(self, row):
613 """ 614 Find the first starred element in the specified row. Returns 615 the column index, or -1 if no starred element was found. 616 """ 617 col = -1 618 for j in range(self.n): 619 if self.marked[row][j] == 1: 620 col = j 621 break 622 623 return col
624
625 - def __find_star_in_col(self, col):
626 """ 627 Find the first starred element in the specified row. Returns 628 the row index, or -1 if no starred element was found. 629 """ 630 row = -1 631 for i in range(self.n): 632 if self.marked[i][col] == 1: 633 row = i 634 break 635 636 return row
637
638 - def __find_prime_in_row(self, row):
639 """ 640 Find the first prime element in the specified row. Returns 641 the column index, or -1 if no starred element was found. 642 """ 643 col = -1 644 for j in range(self.n): 645 if self.marked[row][j] == 2: 646 col = j 647 break 648 649 return col
650
651 - def __convert_path(self, path, count):
652 for i in range(count+1): 653 if self.marked[path[i][0]][path[i][1]] == 1: 654 self.marked[path[i][0]][path[i][1]] = 0 655 else: 656 self.marked[path[i][0]][path[i][1]] = 1
657
658 - def __clear_covers(self):
659 """Clear all covered matrix cells""" 660 for i in range(self.n): 661 self.row_covered[i] = False 662 self.col_covered[i] = False
663
664 - def __erase_primes(self):
665 """Erase all prime markings""" 666 for i in range(self.n): 667 for j in range(self.n): 668 if self.marked[i][j] == 2: 669 self.marked[i][j] = 0
670 671 # --------------------------------------------------------------------------- 672 # Functions 673 # --------------------------------------------------------------------------- 674
675 -def make_cost_matrix(profit_matrix, inversion_function):
676 """ 677 Create a cost matrix from a profit matrix by calling 678 'inversion_function' to invert each value. The inversion 679 function must take one numeric argument (of any type) and return 680 another numeric argument which is presumed to be the cost inverse 681 of the original profit. 682 683 This is a static method. Call it like this: 684 685 .. python:: 686 687 cost_matrix = Munkres.make_cost_matrix(matrix, inversion_func) 688 689 For example: 690 691 .. python:: 692 693 cost_matrix = Munkres.make_cost_matrix(matrix, lambda x : sys.maxint - x) 694 695 :Parameters: 696 profit_matrix : list of lists 697 The matrix to convert from a profit to a cost matrix 698 699 inversion_function : function 700 The function to use to invert each entry in the profit matrix 701 702 :rtype: list of lists 703 :return: The converted matrix 704 """ 705 cost_matrix = [] 706 for row in profit_matrix: 707 cost_matrix.append([inversion_function(value) for value in row]) 708 return cost_matrix
709 742 743 # --------------------------------------------------------------------------- 744 # Main 745 # --------------------------------------------------------------------------- 746 747 if __name__ == '__main__': 748 749 750 matrices = [ 751 # Square 752 ([[400, 150, 400], 753 [400, 450, 600], 754 [300, 225, 300]], 755 850 # expected cost 756 ), 757 758 # Rectangular variant 759 ([[400, 150, 400, 1], 760 [400, 450, 600, 2], 761 [300, 225, 300, 3]], 762 452 # expected cost 763 ), 764 765 # Square 766 ([[10, 10, 8], 767 [ 9, 8, 1], 768 [ 9, 7, 4]], 769 18 770 ), 771 772 # Rectangular variant 773 ([[10, 10, 8, 11], 774 [ 9, 8, 1, 1], 775 [ 9, 7, 4, 10]], 776 15 777 ), 778 ] 779 780 m = Munkres() 781 for cost_matrix, expected_total in matrices: 782 print_matrix(cost_matrix, msg='cost matrix') 783 indexes = m.compute(cost_matrix) 784 total_cost = 0 785 for r, c in indexes: 786 x = cost_matrix[r][c] 787 total_cost += x 788 print '(%d, %d) -> %d' % (r, c, x) 789 print 'lowest cost=%d' % total_cost 790 assert expected_total == total_cost 791