用Python和NumPy实现匈牙利算法_Python_下载.zip

#!/usr/bin/python """ Implementation of the Hungarian (Munkres) Algorithm using Python and NumPy References: http://www.ams.jhu.edu/~castello/362/Handouts/hungarian.pdf http://weber.ucsd.edu/~vcrawfor/hungar.pdf http://en.wikipedia.org/wiki/Hungarian_algorithm http://www.public.iastate.edu/~ddoty/HungarianAlgorithm.html http://www.clapper.org/software/python/munkres/ """ # Module Information. __version__ = "1.1.1" __author__ = "Thom Dedecko" __url__ = "http://github.com/tdedecko/hungarian-algorithm" __copyright__ = "(c) 2010 Thom Dedecko" __license__ = "MIT License" class HungarianError(Exception): pass # Import numpy. Error if fails try: import numpy as np except ImportError: raise HungarianError("NumPy is not installed.") class Hungarian: """ Implementation of the Hungarian (Munkres) Algorithm using np. Usage: hungarian = Hungarian(cost_matrix) hungarian.calculate() or hungarian = Hungarian() hungarian.calculate(cost_matrix) Handle Profit matrix: hungarian = Hungarian(profit_matrix, is_profit_matrix=True) or cost_matrix = Hungarian.make_cost_matrix(profit_matrix) The matrix will be automatically padded if it is not square. For that numpy's resize function is used, which automatically adds 0's to any row/column that is added Get results and total potential after calculation: hungarian.get_results() hungarian.get_total_potential() """ def __init__(self, input_matrix=None, is_profit_matrix=False): """ input_matrix is a List of Lists. input_matrix is assumed to be a cost matrix unless is_profit_matrix is True. """ if input_matrix is not None: # Save input my_matrix = np.array(input_matrix) self._input_matrix = np.array(input_matrix) self._maxColumn = my_matrix.shape[1] self._maxRow = my_matrix.shape[0] # Adds 0s if any columns/rows are added. Otherwise stays unaltered matrix_size = max(self._maxColumn, self._maxRow) pad_columns = matrix_size - self._maxRow pad_rows = matrix_size - self._maxColumn my_matrix = np.pad(my_matrix, ((0,pad_columns),(0,pad_rows)), 'constant', constant_values=(0)) # Convert matrix to profit matrix if necessary if is_profit_matrix: my_matrix = self.make_cost_matrix(my_matrix) self._cost_matrix = my_matrix self._size = len(my_matrix) self._shape = my_matrix.shape # Results from algorithm. self._results = [] self._totalPotential = 0 else: self._cost_matrix = None def get_results(self): """Get results after calculation.""" return self._results def get_total_potential(self): """Returns expected value after calculation.""" return self._totalPotential def calculate(self, input_matrix=None, is_profit_matrix=False): """ Implementation of the Hungarian (Munkres) Algorithm. input_matrix is a List of Lists. input_matrix is assumed to be a cost matrix unless is_profit_matrix is True. """ # Handle invalid and new matrix inputs. if input_matrix is None and self._cost_matrix is None: raise HungarianError("Invalid input") elif input_matrix is not None: self.__init__(input_matrix, is_profit_matrix) result_matrix = self._cost_matrix.copy() # Step 1: Subtract row mins from each row. for index, row in enumerate(result_matrix): result_matrix[index] -= row.min() # Step 2: Subtract column mins from each column. for index, column in enumerate(result_matrix.T): result_matrix[:, index] -= column.min() # Step 3: Use minimum number of lines to cover all zeros in the matrix. # If the total covered rows+columns is not equal to the matrix size then adjust matrix and repeat. total_covered = 0 while total_covered < self._size: # Find minimum number of lines to cover all zeros in the matrix and find total covered rows and columns. cover_zeros = CoverZeros(result_matrix) covered_rows = cover_zeros.get_covered_rows() covered_columns = cover_zeros.get_covered_columns() total_covered = len(covered_rows) + len(covered_columns) # if the total covered rows+columns is not equal to the matrix size then adjust it by min uncovered num (m). if total_covered < self._size: result_matrix = self._adjust_matrix_by_min_uncovered_num(result_matrix, covered_rows, covered_columns) # Step 4: Starting with the top row, work your way downwards as you make assignments. # Find single zeros in rows or columns. # Add them to final result and remove them and their associated row/column from the matrix. expected_results = min(self._maxColumn, self._maxRow) zero_locations = (result_matrix == 0) while len(self._results) != expected_results: # If number of zeros in the matrix is zero before finding all the results then an error has occurred. if not zero_locations.any(): raise HungarianError("Unable to find results. Algorithm has failed.") # Find results and mark rows and columns for deletion matched_rows, matched_columns = self.__find_matches(zero_locations) # Make arbitrary selection total_matched = len(matched_rows) + len(matched_columns) if total_matched == 0: matched_rows, matched_columns = self.select_arbitrary_match(zero_locations) # Delete rows and columns for row in matched_rows: zero_locations[row] = False for column in matched_columns: zero_locations[:, column] = False # Save Results self.__set_results(zip(matched_rows, matched_columns)) # Calculate total potential value = 0 for row, column in self._results: value += self._input_matrix[row, column] self._totalPotential = value @staticmethod def make_cost_matrix(profit_matrix): """ Converts a profit matrix into a cost matrix. Expects NumPy objects as input. """ # subtract profit matrix from a matrix made of the max value of the profit matrix matrix_shape = profit_matrix.shape offset_matrix = np.ones(matrix_shape, dtype=int) * profit_matrix.max() cost_matrix = offset_matrix - profit_matrix return cost_matrix def _adjust_matrix_by_min_uncovered_num(self, result_matrix, covered_rows, covered_columns): """Subtract m from every uncovered number and add m to every element covered with two lines.""" # Calculate minimum uncovered number (m) elements = [] for row_index, row in enumerate(result_matrix): if row_index not in covered_rows: for index, element in enumerate(row): if index not in covered_columns: elements.append(element) min_uncovered_num = min(elements) # Add m to every covered element adjusted_matrix = result_matrix for row in covered_rows: adjusted_matrix[row] += min_uncovered_num for column in covered_columns: adjusted_matrix[:, column] += min_uncovered_num # Subtract m from every element m_matrix = np.ones(self._shape, dtype=int) * min_uncovered_num adjusted_matrix -= m_matrix return adjusted_matrix def __find_matches(self, zero_locations): """Returns rows and columns with matches in them.""" marked_rows = np.array([], dtype=int) marke