Skip to content

Instantly share code, notes, and snippets.

@Gnimuc

Gnimuc/Munkres.jl,mem

Created February 12, 2018 10:05
Show Gist options
  • Save Gnimuc/ac875c62a08ec5a925178a9c04b89a9d to your computer and use it in GitHub Desktop.
Save Gnimuc/ac875c62a08ec5a925178a9c04b89a9d to your computer and use it in GitHub Desktop.
Download ZIP
Updated
Raw
Munkres.jl,mem
- """
- munkres(costMat) -> Zs
-
- Find an optimal solution of the assignment problem represented by the square
- matrix `costMat`. Return an sparse matrix illustrating the optimal matching.
-
- # Examples
-
- ```julia
- julia> costMat = ones(3, 3) - eye(3, 3)
- 3×3 Array{Float64,2}:
- 0.0 1.0 1.0
- 1.0 0.0 1.0
- 1.0 1.0 0.0
-
- julia> matching = Hungarian.munkres(costMat)
- 3×3 sparse matrix with 3 Int64 nonzero entries:
- [1, 1] = 2
- [2, 2] = 2
- [3, 3] = 2
-
- julia> full(matching)
- 3×3 Array{Int64,2}:
- 2 0 0
- 0 2 0
- 0 0 2
- ```
- """
- function munkres(costMat::AbstractMatrix{T}) where T <: Real
0 size(costMat,2) ≥ size(costMat,1) || throw(ArgumentError("Non-square matrix should have more columns than rows."))
-
8008000 A = copy(costMat)
-
- # preliminaries:
- # "no lines are covered;"
14400 rowCovered = falses(size(A,1))
14400 colCovered = falses(size(A,2))
-
- # for tracking changes per row/col of A
89600 Δrow = zeros(size(A,1))
89600 Δcol = zeros(size(A,2))
-
- # for tracking indices
8000 rowCoveredIdx = Int[]
8000 colCoveredIdx = Int[]
8000 rowUncoveredIdx = Int[]
8000 colUncoveredIdx = Int[]
0 sizehint!(rowCoveredIdx, size(A,1))
0 sizehint!(colCoveredIdx, size(A,2))
0 sizehint!(rowUncoveredIdx, size(A,1))
0 sizehint!(colUncoveredIdx, size(A,2))
-
- # "no zeros are starred or primed."
- # use a sparse matrix Zs to store these three kinds of zeros:
- # 0 => NON => Non-zero
- # 1 => Z => ordinary zero
- # 2 => STAR => starred zero
- # 3 => PRIME => primed zero
112000 Zs = spzeros(Int8, size(A)...)
-
- # "consider a row of the matrix A;
- # subtract from each element in this row the smallest element of this row.
- # do the same for each row."
0 A .-= minimum(A, 2)
-
- # "then consider each column of the resulting matrix and subtract from each
- # column its smallest entry."
- # Note that, this step should be removed if the input matrix is not square.
- # A .-= minimum(A, 1)
-
14400 rowSTAR = falses(size(A,1))
14400 colSTAR = falses(size(A,2))
8000 row2colSTAR = Dict{Int,Int}()
0 for ii in CartesianRange(size(A))
- # "consider a zero Z of the matrix;"
0 if A[ii] == 0
259200 Zs[ii] = Z
- # "if there is no starred zero in its row and none in its column, star Z.
- # repeat, considering each zero in the matrix in turn;"
0 r, c = ii.I
0 if !colSTAR[c] && !rowSTAR[r]
0 Zs[r,c] = STAR
0 rowSTAR[r] = true
0 colSTAR[c] = true
616000 row2colSTAR[r] = c
- # "then cover every column containing a starred zero."
0 colCovered[c] = true
- end
- end
- end
-
- # preliminaries done, go to step 1
0 stepNum = 1
-
- # if the assignment is already found, exist
- # here we adjust Munkres's algorithm in order to deal with rectangular matrices,
- # so only K columns are counted here, where K = min(size(Zs))
60800 if length(find(colCovered)) == minimum(size(Zs))
0 stepNum = 0
- end
-
- # these three steps are parallel with those in the paper:
- # J. Munkres, "Algorithms for the Assignment and Transportation Problems",
- # Journal of the Society for Industrial and Applied Mathematics, 5(1):32–38, 1957 March.
0 while stepNum != 0
0 if stepNum == 1
0 stepNum = step1!(Zs, rowCovered, colCovered, rowSTAR, row2colSTAR)
0 elseif stepNum == 2
0 stepNum = step2!(Zs, rowCovered, colCovered, rowSTAR, row2colSTAR)
0 elseif stepNum == 3
0 stepNum = step3!(A, Zs, rowCovered, colCovered, rowSTAR, row2colSTAR, Δrow, Δcol, rowCoveredIdx, colCoveredIdx, rowUncoveredIdx, colUncoveredIdx)
- end
- end
-
0 return Zs
- end
-
- function munkres(costMat::AbstractMatrix{S}) where {T <: Real, S <: Union{Missing, T}}
- # replace forbidden edges (i.e. those with a missing cost) by a very large cost, so that they
- # are never chosen for the matching (except if they are the only possible edges)
- costMatReal = [ismissing(x) ? typemax(T) : x for x in costMat]
-
- # compute the assignment with the standard procedure. the return value is guaranteed to be sparse
- assignment = munkres(costMatReal)
-
- # remove the forbidden edges, would they be in the assignment
- colLen = size(assignment, 2)
- rows = rowvals(assignment)
- for c in 1:colLen, i in nzrange(assignment, c)
- r = rows[i]
- if assignment[r, c] == STAR && ismissing(costMat[r, c])
- assignment[r, c] = Z # standard zero, no assignment made
- end
- end
-
- return assignment
- end
-
- """
- Step 1 of the original Munkres' Assignment Algorithm
- """
- function step1!(Zs, rowCovered, colCovered, rowSTAR, row2colSTAR)
0 colLen = size(Zs,2)
0 rows = rowvals(Zs)
- # step 1:
0 zeroCoveredNum = 0
- # "repeat until all zeros are covered."
0 while zeroCoveredNum < nnz(Zs)
0 zeroCoveredNum = 0
0 for c = 1:colLen, i in nzrange(Zs, c)
0 r = rows[i]
- # "choose a non-covered zero and prime it"
0 if colCovered[c] == false && rowCovered[r] == false
0 Zs[r,c] = PRIME
- # "consider the row containing it."
- # "if there is a starred zero Z in this row"
0 if rowSTAR[r]
- # "cover this row and uncover the column of Z"
0 rowCovered[r] = true
0 colCovered[row2colSTAR[r]] = false
- else
- # "if there is no starred zero in this row,
- # go at once to step 2."
0 return 2
- end
- else
- # otherwise, this zero is covered
0 zeroCoveredNum += 1
- end
- end
- end
- # "go to step 3."
0 return 3
- end
-
- """
- Step 2 of the original Munkres' Assignment Algorithm
- """
- function step2!(Zs, rowCovered, colCovered, rowSTAR, row2colSTAR)
0 ZsDims = size(Zs)
0 rows = rowvals(Zs)
- # step 2:
336000 sequence = Tuple{Int,Int}[]
0 flag = false
- # "there is a sequence of alternating primed and starred zeros, constructed
- # as follows:"
- # "let Z₀ denote the uncovered 0′.[there is only one.]"
0 for c = 1:ZsDims[2], i in nzrange(Zs, c)
0 r = rows[i]
- # note that Z₀ is an **uncovered** 0′
0 if Zs[r,c] == PRIME && colCovered[c] == false && rowCovered[r] == false
- # push Z₀(uncovered 0′) into the sequence
0 push!(sequence, (r, c))
- # "let Z₁ denote the 0* in the Z₀'s column.(if any)"
- # find 0* in Z₀'s column
0 for j in nzrange(Zs, c)
0 Z₁r = rows[j]
0 if Zs[Z₁r, c] == STAR
- # push Z₁(0*) into the sequence
0 push!(sequence, (Z₁r, c))
- # set sequence continue flag => true
0 flag = true
0 break
- end
- end
0 break
- end
- end
- # "let Z₂ denote the 0′ in Z₁'s row(there will always be one)."
- # "let Z₃ denote the 0* in the Z₂'s column."
- # "continue until the sequence stops at a 0′ Z₂ⱼ, which has no 0* in its column."
0 while flag
0 flag = false
0 r = sequence[end][1]
- # find Z₂ in Z₃'s row (always exists)
0 for c = 1:ZsDims[2]
0 if Zs[r,c] == PRIME
- # push Z₂ into the sequence
0 push!(sequence, (r, c))
- # find 0* in Z₂'s column
0 for j in nzrange(Zs, c)
0 Z₃r = rows[j]
0 if Zs[Z₃r, c] == STAR
- # push Z₃(0*) into the sequence
0 push!(sequence, (Z₃r, c))
- # set sequence continue flag => true
0 flag = true
0 break
- end
- end
0 break
- end
- end
- end
-
- # "unstar each starred zero of the sequence;"
0 for i in 2:2:endof(sequence)-1
0 Zs[sequence[i]...] = Z
- end
-
- # clean up
0 fill!(rowSTAR, false)
0 empty!(row2colSTAR)
-
- # "and star each primed zero of the sequence."
0 for i in 1:2:endof(sequence)
0 Zs[sequence[i]...] = STAR
- end
-
0 for c = 1:ZsDims[2], i in nzrange(Zs, c)
0 r = rows[i]
- # "erase all primes;"
0 if Zs[r,c] == PRIME
0 Zs[r,c] = Z
- end
- # "and cover every column containing a 0*."
0 if Zs[r,c] == STAR
0 colCovered[c] = true
0 rowSTAR[r] = true
0 row2colSTAR[r] = c
- end
- end
-
- # "uncover every row"
0 fill!(rowCovered, false)
-
- # "if all columns are covered, the starred zeros form the desired independent set."
- # here we adjust Munkres's algorithm in order to deal with rectangular matrices,
- # so only K columns are counted here, where K = min(size(Zs))
3424000 if length(find(colCovered)) == minimum(size(Zs))
- # algorithm exits
0 return 0
- else
- # "otherwise, return to step 1."
0 return 1
- end
- end
-
- """
- Step 3 of the original Munkres' Assignment Algorithm
- """
- function step3!(A::AbstractMatrix{T}, Zs, rowCovered, colCovered, rowSTAR, row2colSTAR,
- Δrow, Δcol, rowCoveredIdx, colCoveredIdx, rowUncoveredIdx, colUncoveredIdx) where T <: Real
- # step 3(Step C):
- # "let h denote the smallest uncovered element of the matrix;
- # add h to each covered row; then subtract h from each uncovered column."
- # -h ||
- # | no change || reduce old zeros
- # | change:(+h)+(-h)=0 || change: +h
- # +h ---------Covered Row-----||----Covered Row--------
- # | Uncovered Column || Covered Column
- #==============|====================||================================#
- # | Uncovered Row || Uncovered Row
- # | Uncovered Column || Covered Column
- # | change: -h || change: 0
- # | produce new zeros || no change
- # | ||
- # since we always apply add/substract h to a whole row/column, it's unnecessary
- # to apply every operation to every entry of A's row/column, we only need a row
- # and a column vector to keep tracking those changes and use it when necessary.
-
- # find h and track the location of those new zeros
0 empty!(rowCoveredIdx)
0 empty!(colCoveredIdx)
0 empty!(rowUncoveredIdx)
0 empty!(colUncoveredIdx)
-
0 for i in eachindex(rowCovered)
0 rowCovered[i] ? push!(rowCoveredIdx, i) : push!(rowUncoveredIdx, i)
- end
-
0 for i in eachindex(colCovered)
0 colCovered[i] ? push!(colCoveredIdx, i) : push!(colUncoveredIdx, i)
- end
-
0 h = Inf
1120000 minLocations = Tuple{Int,Int}[]
0 @inbounds for j in colUncoveredIdx, i in rowUncoveredIdx
0 cost = A[i,j] + Δcol[j] + Δrow[i]
0 if cost <= h
0 if cost != h
0 h = cost
0 empty!(minLocations)
- end
0 push!(minLocations, (i,j))
- end
- end
-
0 for i in rowCoveredIdx
0 Δrow[i] += h
- end
-
0 for i in colUncoveredIdx
0 Δcol[i] -= h
- end
-
- # mark new zeros in Zs and remove those elements that will no longer be zero
- # produce new zeros
0 for loc in minLocations
244800 Zs[loc...] = Z
- end
- # reduce old zeros
0 rows = rowvals(Zs)
0 for c in colCoveredIdx, i in nzrange(Zs, c)
0 r = rows[i]
0 if rowCovered[r]
0 if Zs[r,c] == STAR
0 rowSTAR[r] = false
0 delete!(row2colSTAR, r)
- end
0 Zs[r,c] = NON
- end
- end
-
0 dropzeros!(Zs)
-
- # "return to step 1."
0 return stepNum = 1
- end
-
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment