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.maxsize 
 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.maxsize - 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.maxsize - 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.maxsize - 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  import copy 
281   
282  # --------------------------------------------------------------------------- 
283  # Exports 
284  # --------------------------------------------------------------------------- 
285   
286  __all__     = ['Munkres', 'make_cost_matrix'] 
287   
288  # --------------------------------------------------------------------------- 
289  # Globals 
290  # --------------------------------------------------------------------------- 
291   
292  # Info about the module 
293  __version__   = "1.0.6" 
294  __author__    = "Brian Clapper, bmc@clapper.org" 
295  __url__       = "http://software.clapper.org/munkres/" 
296  __copyright__ = "(c) 2008 Brian M. Clapper" 
297  __license__   = "BSD-style license" 
298   
299  # --------------------------------------------------------------------------- 
300  # Classes 
301  # --------------------------------------------------------------------------- 
302   
303 -class Munkres:
304 """ 305 Calculate the Munkres solution to the classical assignment problem. 306 See the module documentation for usage. 307 """ 308
309 - def __init__(self):
310 """Create a new instance""" 311 self.C = None 312 self.row_covered = [] 313 self.col_covered = [] 314 self.n = 0 315 self.Z0_r = 0 316 self.Z0_c = 0 317 self.marked = None 318 self.path = None
319
320 - def make_cost_matrix(profit_matrix, inversion_function):
321 """ 322 **DEPRECATED** 323 324 Please use the module function ``make_cost_matrix()``. 325 """ 326 import munkres 327 return munkres.make_cost_matrix(profit_matrix, inversion_function)
328 329 make_cost_matrix = staticmethod(make_cost_matrix) 330
331 - def pad_matrix(self, matrix, pad_value=0):
332 """ 333 Pad a possibly non-square matrix to make it square. 334 335 :Parameters: 336 matrix : list of lists 337 matrix to pad 338 339 pad_value : int 340 value to use to pad the matrix 341 342 :rtype: list of lists 343 :return: a new, possibly padded, matrix 344 """ 345 max_columns = 0 346 total_rows = len(matrix) 347 348 for row in matrix: 349 max_columns = max(max_columns, len(row)) 350 351 total_rows = max(max_columns, total_rows) 352 353 new_matrix = [] 354 for row in matrix: 355 row_len = len(row) 356 new_row = row[:] 357 if total_rows > row_len: 358 # Row too short. Pad it. 359 new_row += [0] * (total_rows - row_len) 360 new_matrix += [new_row] 361 362 while len(new_matrix) < total_rows: 363 new_matrix += [[0] * total_rows] 364 365 return new_matrix
366
367 - def compute(self, cost_matrix):
368 """ 369 Compute the indexes for the lowest-cost pairings between rows and 370 columns in the database. Returns a list of (row, column) tuples 371 that can be used to traverse the matrix. 372 373 :Parameters: 374 cost_matrix : list of lists 375 The cost matrix. If this cost matrix is not square, it 376 will be padded with zeros, via a call to ``pad_matrix()``. 377 (This method does *not* modify the caller's matrix. It 378 operates on a copy of the matrix.) 379 380 **WARNING**: This code handles square and rectangular 381 matrices. It does *not* handle irregular matrices. 382 383 :rtype: list 384 :return: A list of ``(row, column)`` tuples that describe the lowest 385 cost path through the matrix 386 387 """ 388 self.C = self.pad_matrix(cost_matrix) 389 self.n = len(self.C) 390 self.original_length = len(cost_matrix) 391 self.original_width = len(cost_matrix[0]) 392 self.row_covered = [False for i in range(self.n)] 393 self.col_covered = [False for i in range(self.n)] 394 self.Z0_r = 0 395 self.Z0_c = 0 396 self.path = self.__make_matrix(self.n * 2, 0) 397 self.marked = self.__make_matrix(self.n, 0) 398 399 done = False 400 step = 1 401 402 steps = { 1 : self.__step1, 403 2 : self.__step2, 404 3 : self.__step3, 405 4 : self.__step4, 406 5 : self.__step5, 407 6 : self.__step6 } 408 409 while not done: 410 try: 411 func = steps[step] 412 step = func() 413 except KeyError: 414 done = True 415 416 # Look for the starred columns 417 results = [] 418 for i in range(self.original_length): 419 for j in range(self.original_width): 420 if self.marked[i][j] == 1: 421 results += [(i, j)] 422 423 return results
424
425 - def __copy_matrix(self, matrix):
426 """Return an exact copy of the supplied matrix""" 427 return copy.deepcopy(matrix)
428
429 - def __make_matrix(self, n, val):
430 """Create an *n*x*n* matrix, populating it with the specific value.""" 431 matrix = [] 432 for i in range(n): 433 matrix += [[val for j in range(n)]] 434 return matrix
435
436 - def __step1(self):
437 """ 438 For each row of the matrix, find the smallest element and 439 subtract it from every element in its row. Go to Step 2. 440 """ 441 C = self.C 442 n = self.n 443 for i in range(n): 444 minval = min(self.C[i]) 445 # Find the minimum value for this row and subtract that minimum 446 # from every element in the row. 447 for j in range(n): 448 self.C[i][j] -= minval 449 450 return 2
451
452 - def __step2(self):
453 """ 454 Find a zero (Z) in the resulting matrix. If there is no starred 455 zero in its row or column, star Z. Repeat for each element in the 456 matrix. Go to Step 3. 457 """ 458 n = self.n 459 for i in range(n): 460 for j in range(n): 461 if (self.C[i][j] == 0) and \ 462 (not self.col_covered[j]) and \ 463 (not self.row_covered[i]): 464 self.marked[i][j] = 1 465 self.col_covered[j] = True 466 self.row_covered[i] = True 467 468 self.__clear_covers() 469 return 3
470
471 - def __step3(self):
472 """ 473 Cover each column containing a starred zero. If K columns are 474 covered, the starred zeros describe a complete set of unique 475 assignments. In this case, Go to DONE, otherwise, Go to Step 4. 476 """ 477 n = self.n 478 count = 0 479 for i in range(n): 480 for j in range(n): 481 if self.marked[i][j] == 1: 482 self.col_covered[j] = True 483 count += 1 484 485 if count >= n: 486 step = 7 # done 487 else: 488 step = 4 489 490 return step
491
492 - def __step4(self):
493 """ 494 Find a noncovered zero and prime it. If there is no starred zero 495 in the row containing this primed zero, Go to Step 5. Otherwise, 496 cover this row and uncover the column containing the starred 497 zero. Continue in this manner until there are no uncovered zeros 498 left. Save the smallest uncovered value and Go to Step 6. 499 """ 500 step = 0 501 done = False 502 row = -1 503 col = -1 504 star_col = -1 505 while not done: 506 (row, col) = self.__find_a_zero() 507 if row < 0: 508 done = True 509 step = 6 510 else: 511 self.marked[row][col] = 2 512 star_col = self.__find_star_in_row(row) 513 if star_col >= 0: 514 col = star_col 515 self.row_covered[row] = True 516 self.col_covered[col] = False 517 else: 518 done = True 519 self.Z0_r = row 520 self.Z0_c = col 521 step = 5 522 523 return step
524
525 - def __step5(self):
526 """ 527 Construct a series of alternating primed and starred zeros as 528 follows. Let Z0 represent the uncovered primed zero found in Step 4. 529 Let Z1 denote the starred zero in the column of Z0 (if any). 530 Let Z2 denote the primed zero in the row of Z1 (there will always 531 be one). Continue until the series terminates at a primed zero 532 that has no starred zero in its column. Unstar each starred zero 533 of the series, star each primed zero of the series, erase all 534 primes and uncover every line in the matrix. Return to Step 3 535 """ 536 count = 0 537 path = self.path 538 path[count][0] = self.Z0_r 539 path[count][1] = self.Z0_c 540 done = False 541 while not done: 542 row = self.__find_star_in_col(path[count][1]) 543 if row >= 0: 544 count += 1 545 path[count][0] = row 546 path[count][1] = path[count-1][1] 547 else: 548 done = True 549 550 if not done: 551 col = self.__find_prime_in_row(path[count][0]) 552 count += 1 553 path[count][0] = path[count-1][0] 554 path[count][1] = col 555 556 self.__convert_path(path, count) 557 self.__clear_covers() 558 self.__erase_primes() 559 return 3
560
561 - def __step6(self):
562 """ 563 Add the value found in Step 4 to every element of each covered 564 row, and subtract it from every element of each uncovered column. 565 Return to Step 4 without altering any stars, primes, or covered 566 lines. 567 """ 568 minval = self.__find_smallest() 569 for i in range(self.n): 570 for j in range(self.n): 571 if self.row_covered[i]: 572 self.C[i][j] += minval 573 if not self.col_covered[j]: 574 self.C[i][j] -= minval 575 return 4
576
577 - def __find_smallest(self):
578 """Find the smallest uncovered value in the matrix.""" 579 minval = sys.maxsize 580 for i in range(self.n): 581 for j in range(self.n): 582 if (not self.row_covered[i]) and (not self.col_covered[j]): 583 if minval > self.C[i][j]: 584 minval = self.C[i][j] 585 return minval
586
587 - def __find_a_zero(self):
588 """Find the first uncovered element with value 0""" 589 row = -1 590 col = -1 591 i = 0 592 n = self.n 593 done = False 594 595 while not done: 596 j = 0 597 while True: 598 if (self.C[i][j] == 0) and \ 599 (not self.row_covered[i]) and \ 600 (not self.col_covered[j]): 601 row = i 602 col = j 603 done = True 604 j += 1 605 if j >= n: 606 break 607 i += 1 608 if i >= n: 609 done = True 610 611 return (row, col)
612
613 - def __find_star_in_row(self, row):
614 """ 615 Find the first starred element in the specified row. Returns 616 the column index, or -1 if no starred element was found. 617 """ 618 col = -1 619 for j in range(self.n): 620 if self.marked[row][j] == 1: 621 col = j 622 break 623 624 return col
625
626 - def __find_star_in_col(self, col):
627 """ 628 Find the first starred element in the specified row. Returns 629 the row index, or -1 if no starred element was found. 630 """ 631 row = -1 632 for i in range(self.n): 633 if self.marked[i][col] == 1: 634 row = i 635 break 636 637 return row
638
639 - def __find_prime_in_row(self, row):
640 """ 641 Find the first prime element in the specified row. Returns 642 the column index, or -1 if no starred element was found. 643 """ 644 col = -1 645 for j in range(self.n): 646 if self.marked[row][j] == 2: 647 col = j 648 break 649 650 return col
651
652 - def __convert_path(self, path, count):
653 for i in range(count+1): 654 if self.marked[path[i][0]][path[i][1]] == 1: 655 self.marked[path[i][0]][path[i][1]] = 0 656 else: 657 self.marked[path[i][0]][path[i][1]] = 1
658
659 - def __clear_covers(self):
660 """Clear all covered matrix cells""" 661 for i in range(self.n): 662 self.row_covered[i] = False 663 self.col_covered[i] = False
664
665 - def __erase_primes(self):
666 """Erase all prime markings""" 667 for i in range(self.n): 668 for j in range(self.n): 669 if self.marked[i][j] == 2: 670 self.marked[i][j] = 0
671 672 # --------------------------------------------------------------------------- 673 # Functions 674 # --------------------------------------------------------------------------- 675
676 -def make_cost_matrix(profit_matrix, inversion_function):
677 """ 678 Create a cost matrix from a profit matrix by calling 679 'inversion_function' to invert each value. The inversion 680 function must take one numeric argument (of any type) and return 681 another numeric argument which is presumed to be the cost inverse 682 of the original profit. 683 684 This is a static method. Call it like this: 685 686 .. python:: 687 688 cost_matrix = Munkres.make_cost_matrix(matrix, inversion_func) 689 690 For example: 691 692 .. python:: 693 694 cost_matrix = Munkres.make_cost_matrix(matrix, lambda x : sys.maxsize - x) 695 696 :Parameters: 697 profit_matrix : list of lists 698 The matrix to convert from a profit to a cost matrix 699 700 inversion_function : function 701 The function to use to invert each entry in the profit matrix 702 703 :rtype: list of lists 704 :return: The converted matrix 705 """ 706 cost_matrix = [] 707 for row in profit_matrix: 708 cost_matrix.append([inversion_function(value) for value in row]) 709 return cost_matrix
710 743 744 # --------------------------------------------------------------------------- 745 # Main 746 # --------------------------------------------------------------------------- 747 748 if __name__ == '__main__': 749 750 matrices = [ 751 # Square 752 ([[400, 150, 400], 753 [400, 450, 600], 754 [300, 225, 300]], 755 850), # expected cost 756 757 # Rectangular variant 758 ([[400, 150, 400, 1], 759 [400, 450, 600, 2], 760 [300, 225, 300, 3]], 761 452), # expected cost 762 763 764 # Square 765 ([[10, 10, 8], 766 [9, 8, 1], 767 [9, 7, 4]], 768 18), 769 770 # Rectangular variant 771 ([[10, 10, 8, 11], 772 [9, 8, 1, 1], 773 [9, 7, 4, 10]], 774 15)] 775 776 m = Munkres() 777 for cost_matrix, expected_total in matrices: 778 print_matrix(cost_matrix, msg='cost matrix') 779 indexes = m.compute(cost_matrix) 780 total_cost = 0 781 for r, c in indexes: 782 x = cost_matrix[r][c] 783 total_cost += x 784 print('(%d, %d) -> %d' % (r, c, x)) 785 print('lowest cost=%d' % total_cost) 786 assert expected_total == total_cost 787