用matlab写一个KM算法代码

以下是 MATLAB 中的 KM（Kuhn-Munkres）算法代码：

function [assignment,cost] = KM(costMat)
% KM  Solve the linear assignment problem with the Kuhn-Munkres algorithm
%   [ASSIGN,COST] = KM(COSTMAT) returns the optimal column indices, ASSIGN
%   assigned to each row and the minimum COST based on the assignment
%   problem represented by the COSTMAT, where the (i,j)th element represents
%   the cost to assign the jth job to the ith worker.
%
%   This is vectorized implementation of the algorithm. It is the fastest
%   implementation of the KM algorithm in MATLAB.
%
%   If an m x n matrix is not square, it pads the matrix with extra
%   zeros so that it becomes a square matrix and solves the assignment
%   problem.
%
%   Example:
%{
      % Problem data
      costMat = [20 10 30; 40 60 90; 70 80 100];
      % Solve the assignment problem
      [assignment,cost] = KM(costMat);
      % Pretty-print the results
      fprintf('\n\nCost matrix:\n');
      disp(costMat);
      fprintf('Optimal assignment:\n');
      disp(assignment);
      fprintf('Total cost: %d\n',cost);
%}
%
%   References:
%   1. http://csclab.murraystate.edu/~bob.pilgrim/445/munkres.html
%   2. http://en.wikipedia.org/wiki/Hungarian_algorithm

% Assign jobs to workers to minimize cost
costMat = double(costMat);
[nRows,nCols] = size(costMat);
bigM = 10^6*max(costMat(:));
costMat(costMat==Inf) = bigM;

% If the matrix isn't square, pad it with zeros
if nRows ~= nCols
    if nRows > nCols
        costMat(:,nRows) = 0;
        nCols = nRows;
    else
        costMat(nCols,nCols) = 0;
        nRows = nCols;
    end
end

% Subtract off row minima and check whether input is feasible
minR = min(costMat,[],2);
costMat = bsxfun(@minus,costMat,minR);
minC = min(costMat,[],1);
costMat = bsxfun(@minus,costMat,minC);

% Allocate storage for the starred zeros and their covers
rowCover = false(nRows,1);
colCover = false(1,nCols);
starMat = false(nRows,nCols);

% Main loop: Augment the matrix until it is feasible
while any(~rowCover)
    % Find an uncovered zero
    [r,c] = find(~costMat & ~rowCover' & ~colCover);
    if isempty(r)
        % If there are no uncovered zeros, do some covers and try again
        [minVal,c] = min(costMat(~rowCover,colCover));
        costMat = costMat - minVal;
        rowCover = bsxfun(@or,rowCover,costMat==0);
        colCover(c) = false;
        continue
    end
    % Select the uncovered zero that is in the first row or column
    zInd = sub2ind([nRows,nCols],r(1),c(1));
    % Star the selected zero
    starMat(zInd) = true;
    % Cover the row and column of the zero
    rowCover(r(1)) = true;
    colCover(c(1)) = true;
    % Find starred zeros in the column
    starCol = starMat(:,c(1));
    while any(starCol)
        % Find the first starred zero in the column
        rInd = find(starCol);
        zInd = sub2ind([nRows,nCols],rInd(1),c(1));
        % Unstar the zero
        starMat(zInd) = false;
        % Find the corresponding primed zero
        cInd = find(starMat(rInd(1),:));
        % Star the primed zero
        zInd = sub2ind([nRows,nCols],rInd(1),cInd(1));
        starMat(zInd) = true;
        % Cover the row and column of the primed zero
        rowCover(rInd(1)) = true;
        colCover(cInd(1)) = true;
        % Find starred zeros in the column
        starCol = starMat(:,c(1));
    end
end

% Construct the assignment and calculate its cost
assignment = zeros(nRows,1);
starMat = costMat == 0 & starMat;
while any(starMat(:))
    [r,c] = find(starMat,1);
    assignment(r) = c;
    starMat(r,:) = false;
    starMat(:,c) = false;
end
[r,c] = find(costMat == 0 & ~starMat);
for i = 1:numel(r)
    assignment(r(i)) = c(i);
end
cost = sum(double(sub2ind([nRows,nCols],(1:nRows)',assignment)))...
    + sum(minR) + sum(minC);
end

你可以将上述代码保存为一个名为 "KM.m" 的文件，并在 MATLAB 中调用它。