munkres

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='Lowest cost 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  Copyright (c) 2008 Brian M. Clapper
    241
    242  This is free software, released under the following BSD-like license:
    243
    244  Redistribution and use in source and binary forms, with or without
    245  modification, are permitted provided that the following conditions are met:
    246
    247  1. Redistributions of source code must retain the above copyright notice,
    248     this list of conditions and the following disclaimer.
    249
    250  2. The end-user documentation included with the redistribution, if any,
    251     must include the following acknowlegement:
    252
    253     This product includes software developed by Brian M. Clapper
    254     (bmc@clapper.org, http://www.clapper.org/bmc/). That software is
    255     copyright (c) 2008 Brian M. Clapper.
    256
    257     Alternately, this acknowlegement may appear in the software itself, if
    258     and wherever such third-party acknowlegements normally appear.
    259
    260  THIS SOFTWARE IS PROVIDED **AS IS**, AND ANY EXPRESSED OR IMPLIED
    261  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
    262  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
    263  EVENT SHALL BRIAN M. CLAPPER BE LIABLE FOR ANY DIRECT, INDIRECT,
    264  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
    265  NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
    266  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
    267  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    268  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
    269  THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    270
    271  ``$Id$``
    272  """
    273
    274  __docformat__ = 'restructuredtext'
    275
    276  # ---------------------------------------------------------------------------
    277  # Imports
    278  # ---------------------------------------------------------------------------
    279
    280  import sys
    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.5.3"
    294  __author__    = "Brian Clapper, bmc@clapper.org"
    295  __url__       = "http://www.clapper.org/software/python/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.maxint
    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.maxint - 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
    711 -def print_matrix(matrix, msg=None):
    712      """
    713      Convenience function: Displays the contents of a matrix of integers.
    714
    715      :Parameters:
    716          matrix : list of lists
    717              Matrix to print
    718
    719          msg : str
    720              Optional message to print before displaying the matrix
    721      """
    722      import math
    723
    724      if msg is not None:
    725          print msg
    726
    727      # Calculate the appropriate format width.
    728      width = 0
    729      for row in matrix:
    730          for val in row:
    731              width = max(width, int(math.log10(val)) + 1)
    732
    733      # Make the format string
    734      format = '%%%dd' % width
    735
    736      # Print the matrix
    737      for row in matrix:
    738          sep = '['
    739          for val in row:
    740              sys.stdout.write(sep + format % val)
    741              sep = ', '
    742          sys.stdout.write(']\n')
    743
    744  # ---------------------------------------------------------------------------
    745  # Main
    746  # ---------------------------------------------------------------------------
    747
    748  if __name__ == '__main__':
    749
    750
    751      matrices = [
    752                  # Square
    753                  ([[400, 150, 400],
    754                    [400, 450, 600],
    755                    [300, 225, 300]],
    756                   850 # expected cost
    757                  ),
    758
    759                  # Rectangular variant
    760                  ([[400, 150, 400, 1],
    761                    [400, 450, 600, 2],
    762                    [300, 225, 300, 3]],
    763                   452 # expected cost
    764                  ),
    765
    766                  # Square
    767                  ([[10, 10,  8],
    768                    [ 9,  8,  1],
    769                    [ 9,  7,  4]],
    770                   18
    771                  ),
    772
    773                  # Rectangular variant
    774                  ([[10, 10,  8, 11],
    775                    [ 9,  8,  1, 1],
    776                    [ 9,  7,  4, 10]],
    777                   15
    778                  ),
    779                 ]
    780
    781      m = Munkres()
    782      for cost_matrix, expected_total in matrices:
    783          print_matrix(cost_matrix, msg='cost matrix')
    784          indexes = m.compute(cost_matrix)
    785          total_cost = 0
    786          for r, c in indexes:
    787              x = cost_matrix[r][c]
    788              total_cost += x
    789              print '(%d, %d) -> %d' % (r, c, x)
    790          print 'lowest cost=%d' % total_cost
    791          assert expected_total == total_cost
    792