docs: def lapjv(i, j, costs, wants_dual_variables = False, augmenting_row_reductions = 2):
      '''Sparse linear assignment solution using Jonker-Volgenant algorithm
      
      i,j - similarly-sized vectors that pair the object at index i[n] with
            the object at index j[j]
            
      costs - a vector of similar size to i and j that is the cost of pairing
              i[n] with j[n].
              
      wants_dual_variables - the dual problem reduces the costs using two
              vectors, u[i] and v[j] where the solution is the maximum value of
              np.sum(u) + np.sum(v) where cost[i,j] - u[i] - v[j] >=  0.
              Set wants_dual_variables to True to have u and v returned in
              addition to the assignments.
              
      augmenting_row_reductions - the authors suggest that augmenting row reduction
              be performed twice to optimize the u and v before the augmenting
              stage. The caller can choose a different number of reductions
              by supplying a different value here.
  
      All costs not appearing in i,j are taken as infinite. Each i in the range,
      0 to max(i) must appear at least once and similarly for j.
      
      returns (x, y), the pairs of assignments that represent the solution
      or (x, y, u, v) if the dual variables are requested.
      '''
      import os
      i = np.atleast_1d(i).astype(int)
      j = np.atleast_1d(j).astype(int)
      costs = np.atleast_1d(costs)
      
      assert len(i) == len(j), "i and j must be the same length"
      assert len(i) == len(costs), "costs must be the same length as i"
  
      #
      # Find the number of i with non-infinite cost for each j
      #
      j_count = np.bincount(j)
      assert not np.any(j_count == 0), "all j must be paired with at least one i"
      #
      # if you order anything by j, this is an index to the minimum for each j
      #
      j_index = np.hstack([[0], np.cumsum(j_count[:-1])])
      #
      # Likewise for i
      #
      i_count = np.bincount(i)
      assert not np.any(i_count == 0), "all i must be paired with at least one j"
      i_index = np.hstack([[0], np.cumsum(i_count[:-1])])
      
      n = len(j_count) # dimension of the square cost matrix
      assert n == len(i_count), "There must be the same number of unique i and j"
      
      # # # # # # # #
      #
      # Variable initialization:
      #
      # The output variables:
      #
      # x - for each i, the assigned j. -1 indicates uninitialized
      # y - for each j, the assigned i
      # u, v - the dual variables
      #
      # A value of x = n or y = n means "unassigned"
      #
      x = np.ascontiguousarray(np.ones(n, np.uint32) * n)
      y = np.ascontiguousarray(np.ones(n, np.uint32) * n, np.uint32)
      u = np.ascontiguousarray(np.zeros(n, np.float64))
      
      # # # # # # # #
      #
      # Column reduction
      #
      # # # # # # # #
      #
      # For a given j, find the i with the minimum cost.
      #
      order = np.lexsort((-i, costs, j))
      min_idx = order[j_index]
      min_i = i[min_idx]
      #
      # v[j] is assigned to the minimum cost over all i
      #
      v = np.ascontiguousarray(costs[min_idx], np.float64)
      #
      # Find the last j for which i was min_i.
      #
      x[min_i] = np.arange(n).astype(np.uint32)
      y[x[x != n]] = np.arange(n).astype(np.uint32)[x != n]
      #
      # Three cases for i:
      #
      # i is not the minimum of any j - i goes on free list
      # i is the minimum of one j - v[j] remains the same and y[x[j]] = i
      # i is the minimum of more than one j, perform reduction transfer
      #
      assignment_count = np.bincount(min_i[min_i != n])
      assignment_count = np.hstack(
          (assignment_count, np.zeros(n - len(assignment_count), int)))
      free_i = assignment_count == 0
      one_i = assignment_count == 1
      #order = np.lexsort((costs, i)) Replace with this after all is done
      order = np.lexsort((j,i))
      j = np.ascontiguousarray(j[order], np.uint32)
      costs = np.ascontiguousarray(costs[order], np.float64)
      i_index = np.ascontiguousarray(i_index, np.uint32)
      i_count = np.ascontiguousarray(i_count, np.uint32)
      if np.any(one_i): 
          reduction_transfer(
              np.ascontiguousarray(np.argwhere(one_i).flatten(), np.uint32),
              j, i_index, i_count, x, u, v, costs)
      #
      # Perform augmenting row reduction on unassigned i
      #
      ii = np.ascontiguousarray(np.argwhere(free_i).flatten(), np.uint32)
      if len(ii) > 0:
          for iii in range (augmenting_row_reductions):
              ii = augmenting_row_reduction(
                  n, ii, j, i_index, i_count, x, y, u, v, costs)
      augment(n, ii,
              j, i_index, i_count, x, y, u, v, costs)
      if wants_dual_variables:
          return x,y,u,v
      else:
          return x,y