Using lapmod with non-square matrices

lapmod.py

def prepare_sparse_cost(shape, cc, ii, jj, cost_limit):
    '''
    Transform the given sparse matrix extending it to a square sparse matrix.
    
    Parameters
    ==========
    shape: tuple
       - cost matrix shape
    (cc, ii, jj): tuple of floats, ints, ints)
        - cost matrix in COO format, see [1].
    cost_limit: float
    
    Returns
    =======
    cc, ii, kk
      - extended square cost matrix in CSR format
    
    Notes
    =====
    WARNING: Avoid using scipy.sparse.coo_matrix(cost) as it will not return the correct (cc, ii, jj).
    `coo_matrix` leaves out any zero values which are the most salient parts of the cost matrix.
    (cc, ii, jj) should include zero costs (if any) and skip all costs that are too large (infinite).
    
    1. https://en.wikipedia.org/wiki/Sparse_matrix
    '''
    assert cost_limit < np.inf
    n, m = shape
    cc_ = np.r_[cc, [cost_limit] * n,
                [cost_limit] * m, [0] * len(cc)]
    ii_ = np.r_[ii, np.arange(0, n, dtype=np.uint32),
                np.arange(n, n + m, dtype=np.uint32), n + jj]
    jj_ = np.r_[jj, np.arange(m, n + m, dtype=np.uint32),
                np.arange(0, m, dtype=np.uint32), m + ii]
    order = np.lexsort((jj_, ii_))
    cc_ = cc_[order]
    kk_ = jj_[order]
    ii_ = np.bincount(ii_, minlength=shape[0]-1)
    ii_ = np.r_[[0], np.cumsum(ii_)]
    ii_ = ii_.astype(np.uint32)
    assert ii_[-1] == 2 * len(cc) + n + m
    return cc_, ii_, kk_
    
cc, ii, kk = prepare_sparse_cost(shape, cc, ii, jj, cost_limit)
ind1, ind0 = lapmod(len(ii)-1, cc, ii, kk, return_cost=False)
ind1[ind1 >= shape[1]] = -1
ind0[ind0 >= shape[0]] = -1
ind1 = ind1[:shape[0]]
ind0 = ind0[:shape[1]]