O(n^3) minimum cost bipartite matching algorithm

bpm.py

# I found all over the place for a O(n^3) minimum cost bipartite matching
# or equivalently maximum weight bipartite matching algorithm but I
# cannot find a clean but precise explanation; the Wikipedia ones,
# Stack Overflow, and course materials are either too high-level
# or only O(n^4).
# So I decided to read the original paper of Edmunds and Karp.  The
# idea is surprisingly simple: just always find a shortest path (that
# has minimum cost) to augment and you will be good.  In fact they
# talked about minimum cost maximum flow problem, but minimum weight
# bipartite matching is just a special case of that.
# Note that this code is only for educational purpose.  It is certainly
# not optimized.

# Weight matrix
D = np.random.rand(10, 10) * 10

N = D.shape[0]
u = np.zeros(N)      # node potentials on LHS
v = np.zeros(N)      # node potentials on RHS
M = np.ones((N, N))  # assignment matrix
x = np.full(N, -1, dtype='int64')    # index of RHS to assign to every LHS node
y = np.full(N, -1, dtype='int64')    # index of LHS to assign to every RHS node
d = D                # the modified weighted matrix of the residual network to find shortest path

for _ in range(N):
    sv = np.full(N, np.inf)          # shortest path from the virtual source to all RHS nodes
    su = np.zeros(N)                 # shortest path from the virtual source to all LHS nodes
    su[x != -1] = np.inf             # if a LHS node is already assigned the edge goes *towards* source instead

    ssu = np.zeros(N, dtype='bool')  # whether a LHS node is visited
    ssv = np.zeros(N, dtype='bool')  # whether a RHS node is visited
    pu = np.full(N, -1, dtype='int64')       # the shortest path predecessor of every LHS as index of RHS
    pv = np.full(N, -1, dtype='int64')       # the shortest path predecessor of every RHS as index of LHS
    s = np.inf       # shortest path cost towards the virtual sink
    pend = -1        # shortest path predecessor towards the virtual sink
    
    # Dijkstra's algorithm to compute shortest path for all nodes
    for _ in range(2 * N):           # for every node in both LHS and RHS
        # Determine the node whose successors' shortest path is to be updated.
        # Could replace with a priority queue but I'm lazy.
        m = np.inf
        xside = True # True means LHS
        mi = 0       # the node ID to visit
        for i in range(N):
            if not ssu[i] and su[i] < m:
                mi = i
                xside = True
                m = su[i]
        for i in range(N):
            if not ssv[i] and sv[i] < m:
                mi = i
                xside = False
                m = sv[i]

        if xside:    # if on LHS...
            for i in range(N):
                # Update all RHS nodes that are not visited already and not assigned to it.
                if i != x[mi] and not ssv[i]:
                    if sv[i] > su[mi] + d[mi, i]:
                        sv[i] = su[mi] + d[mi, i]
                        pv[i] = mi
            ssu[mi] = True
        else:        # if on RHS...
            # If an RHS node is assigned, it must have one and only one successor in the residual network,
            # which is the LHS node assigned to it.  So we update the shortest path of that single LHS
            # node only.
            if y[mi] != -1:
                if su[y[mi]] > sv[mi] + d[y[mi], mi] and not ssu[y[mi]]:
                    su[y[mi]] = sv[mi] + d[y[mi], mi]
                    pu[y[mi]] = mi
            else:    # otherwise update the virtual sink
                if s > sv[mi]:
                    s = sv[mi]
                    pend = mi
            ssv[mi] = True
            
    # Augment along the shortest path.
    while pend != -1:
        xxi = pv[pend]
        yyi = pend
        pend = x[xxi]
        if pend != -1:
            M[xxi, pend] = 1
        x[xxi] = yyi
        y[yyi] = xxi
        M[xxi, yyi] = -1
        
    # Update the weight matrix of the residual network
    u += su
    v += sv
    d = M * (u[:, None] + D - v[None, :])
    
# Solution
print(x)