"""
N-dimensional assignment algorithms (4D and higher).
This module extends the 3D assignment solver to arbitrary dimensions,
enabling more complex assignment scenarios such as:
- 4D: Measurements × Tracks × Hypotheses × Sensors
- 5D+: Additional dimensions for time frames, maneuver classes, etc.
The module provides a unified interface for solving high-dimensional
assignment problems using generalized relaxation methods.
Performance Notes
-----------------
For sparse cost tensors (mostly invalid assignments), use SparseCostTensor
to reduce memory usage by up to 50% and improve performance on large problems.
References
----------
- Poore, A. B., "Multidimensional Assignment Problem and Data
Association," IEEE Transactions on Aerospace and Electronic Systems,
2013.
- Cramer, R. D., et al., "The Emerging Role of Chemical Similarity in
Drug Discovery," Perspectives in Drug Discovery and Design, 2003.
"""
from typing import List, NamedTuple, Optional, Tuple, Union
import numpy as np
from numpy.typing import NDArray
from scipy.optimize import linear_sum_assignment
[docs]
class AssignmentNDResult(NamedTuple):
"""Result of an N-dimensional assignment problem.
Attributes
----------
assignments : ndarray
Array of shape (n_assignments, n_dimensions) containing assigned
index tuples. Each row is an n-tuple of indices.
cost : float
Total assignment cost.
converged : bool
Whether the algorithm converged (for iterative methods).
n_iterations : int
Number of iterations used (for iterative methods).
gap : float
Optimality gap (upper_bound - lower_bound) for relaxation methods.
"""
assignments: NDArray[np.intp]
cost: float
converged: bool
n_iterations: int
gap: float
[docs]
def validate_cost_tensor(cost_tensor: NDArray[np.float64]) -> Tuple[int, ...]:
"""
Validate cost tensor and return dimensions.
Parameters
----------
cost_tensor : ndarray
Cost tensor of arbitrary dimension.
Returns
-------
dims : tuple
Dimensions of the cost tensor.
Raises
------
ValueError
If tensor has fewer than 2 dimensions.
Examples
--------
>>> import numpy as np
>>> cost = np.random.rand(3, 4, 5)
>>> dims = validate_cost_tensor(cost)
>>> dims
(3, 4, 5)
>>> # 1D tensor should raise error
>>> try:
... validate_cost_tensor(np.array([1, 2, 3]))
... except ValueError:
... print("Caught expected error")
Caught expected error
"""
if cost_tensor.ndim < 2:
raise ValueError(
f"Cost tensor must have at least 2 dimensions, got {cost_tensor.ndim}"
)
return cost_tensor.shape
[docs]
def greedy_assignment_nd(
cost_tensor: NDArray[np.float64],
max_assignments: Optional[int] = None,
) -> AssignmentNDResult:
"""
Greedy solver for N-dimensional assignment.
Selects minimum-cost tuples in order until no more valid assignments
exist (no dimension index is repeated).
Parameters
----------
cost_tensor : ndarray
Cost tensor of shape (n1, n2, ..., nk).
max_assignments : int, optional
Maximum number of assignments to find (default: min(dimensions)).
Returns
-------
AssignmentNDResult
Assignments, total cost, and algorithm info.
Examples
--------
>>> import numpy as np
>>> # 3D cost tensor: 3 measurements x 2 tracks x 2 hypotheses
>>> cost = np.array([
... [[1.0, 5.0], [3.0, 2.0]], # meas 0
... [[4.0, 1.0], [2.0, 6.0]], # meas 1
... [[2.0, 3.0], [5.0, 1.0]], # meas 2
... ])
>>> result = greedy_assignment_nd(cost)
>>> result.cost # Total cost of greedy solution
2.0
>>> len(result.assignments) # Number of assignments made
2
Notes
-----
Greedy assignment is fast O(n log n) but not optimal. Used as
heuristic or starting solution for optimization methods.
"""
dims = cost_tensor.shape
n_dims = len(dims)
if max_assignments is None:
max_assignments = min(dims)
# Flatten tensor with index mapping
flat_costs = cost_tensor.ravel()
sorted_indices = np.argsort(flat_costs)
assignments: list[tuple[int, ...]] = []
used_indices: list[set[int]] = [set() for _ in range(n_dims)]
for flat_idx in sorted_indices:
if len(assignments) >= max_assignments:
break
# Convert flat index to multi-dimensional index
multi_idx = np.unravel_index(flat_idx, dims)
# Check if any dimension index is already used
conflict = False
for d, idx in enumerate(multi_idx):
if idx in used_indices[d]:
conflict = True
break
if not conflict:
assignments.append(multi_idx)
for d, idx in enumerate(multi_idx):
used_indices[d].add(idx)
assignments_array = np.array(assignments, dtype=np.intp)
if assignments_array.size > 0:
total_cost = float(np.sum(cost_tensor[tuple(assignments_array.T)]))
else:
total_cost = 0.0
return AssignmentNDResult(
assignments=assignments_array,
cost=total_cost,
converged=True,
n_iterations=1,
gap=0.0, # Greedy doesn't compute lower bound
)
def _recover_feasible_nd(
cost_tensor: NDArray[np.float64],
rows: NDArray[np.intp],
cols: NDArray[np.intp],
) -> Tuple[NDArray[np.intp], float]:
"""Build a feasible N-D assignment from a fixed set of (dim0, dim1) pairs.
The relaxed solution satisfies the first two dimensions' constraints but may
reuse indices in the relaxed dimensions. Each remaining dimension is
assigned by an exact 2-D assignment over the fixed pairs, which guarantees
feasibility and therefore a valid upper bound.
Parameters
----------
cost_tensor : ndarray
Original cost tensor.
rows, cols : ndarray
Indices selected in the first two dimensions.
Returns
-------
assignments : ndarray
Feasible assignment tuples, shape (n_assignments, n_dims).
cost : float
Total cost of the recovered assignment.
"""
dims = cost_tensor.shape
n_dims = len(dims)
n_assign = len(rows)
# Partial tuples: start from the fixed (dim0, dim1) pairs.
partial = [[int(rows[t]), int(cols[t])] for t in range(n_assign)]
for d in range(2, n_dims):
# Cost of extending each partial tuple with each index of dimension d,
# minimized over the not-yet-assigned dimensions.
block = np.empty((n_assign, dims[d]))
for t in range(n_assign):
sub = cost_tensor[tuple(partial[t])]
# sub has axes (dims[d], dims[d+1], ...); minimize the trailing ones
block[t] = sub.reshape(dims[d], -1).min(axis=1)
r_idx, c_idx = linear_sum_assignment(block)
chosen = {int(r): int(c) for r, c in zip(r_idx, c_idx)}
for t in range(n_assign):
# Dimension d may be smaller than the number of tuples; unmatched
# tuples keep the greedy best index (feasibility is enforced only
# among the matched ones, mirroring rectangular 2-D assignment).
partial[t].append(chosen.get(t, int(np.argmin(block[t]))))
assignments = np.array(partial, dtype=np.intp)
if len(assignments) == 0:
return np.empty((0, n_dims), dtype=np.intp), 0.0
cost = float(np.sum(cost_tensor[tuple(assignments.T)]))
return assignments, cost
[docs]
def relaxation_assignment_nd(
cost_tensor: NDArray[np.float64],
max_iterations: int = 100,
tolerance: float = 1e-6,
verbose: bool = False,
) -> AssignmentNDResult:
"""
Lagrangian relaxation solver for N-dimensional assignment.
Relaxes the constraints on dimensions 3..N and solves the resulting
two-dimensional problem exactly, which yields a valid lower bound on the
optimal cost. Subgradient ascent on the multipliers tightens the bound;
a feasible solution is recovered at each iteration to provide an upper
bound.
Parameters
----------
cost_tensor : ndarray
Cost tensor of shape (n1, n2, ..., nk).
max_iterations : int, optional
Maximum iterations (default 100).
tolerance : float, optional
Convergence tolerance for the optimality gap (default 1e-6).
verbose : bool, optional
Print iteration info (default False).
Returns
-------
AssignmentNDResult
Assignments, total cost, convergence info, and optimality gap.
``gap`` is ``best_upper_bound - best_lower_bound``; when ``converged``
is True the returned assignment is provably within ``tolerance`` of
optimal.
Examples
--------
>>> import numpy as np
>>> # 3x3x3 assignment problem
>>> rng = np.random.default_rng(42)
>>> cost = rng.random((3, 3, 3))
>>> result = relaxation_assignment_nd(cost, max_iterations=50)
>>> bool(result.gap >= -1e-9) # gap is a genuine bound
True
>>> result.assignments.shape[1] # 3D assignments
3
Notes
-----
The relaxation follows Poore's formulation (see the module references):
1. Relax the constraints on dimensions 3..N with multipliers ``lambda``.
2. The inner problem separates: minimizing over the relaxed dimensions for
each (i, j) pair leaves a 2-D assignment problem, solved **exactly**.
Solving it exactly is what makes ``L(lambda)`` a valid lower bound --
a greedy inner solve does not bound the optimum and can certify
optimality for suboptimal answers.
3. Recover a feasible solution (upper bound) by assigning each relaxed
dimension with an exact 2-D assignment.
4. Ascend the multipliers along the constraint-violation subgradient.
The reported ``gap`` uses the best bounds seen across all iterations, so
it is a genuine optimality certificate rather than a heuristic score.
"""
dims = validate_cost_tensor(cost_tensor)
n_dims = len(dims)
# A 2-D problem needs no relaxation: solve it exactly.
if n_dims == 2:
rows, cols = linear_sum_assignment(cost_tensor)
assignments = np.column_stack([rows, cols]).astype(np.intp)
cost = float(cost_tensor[rows, cols].sum())
return AssignmentNDResult(
assignments=assignments,
cost=cost,
converged=True,
n_iterations=1,
gap=0.0,
)
free_dims = list(range(2, n_dims))
lambdas = [np.zeros(dims[d]) for d in free_dims]
free_shape = tuple(dims[d] for d in free_dims)
best_upper = np.inf
best_lower = -np.inf
best_assignments: Optional[NDArray[np.intp]] = None
iteration = 0
for iteration in range(max_iterations):
# Relaxed costs: c - sum of multipliers over the relaxed dimensions.
relaxed = cost_tensor.copy()
for k, d in enumerate(free_dims):
shape = [1] * n_dims
shape[d] = dims[d]
relaxed = relaxed - lambdas[k].reshape(shape)
# Minimize over the relaxed dimensions, leaving a 2-D problem.
collapsed = relaxed.reshape(dims[0], dims[1], -1)
reduced = collapsed.min(axis=2)
best_free = collapsed.argmin(axis=2)
# Exact 2-D solve => valid lower bound.
rows, cols = linear_sum_assignment(reduced)
lower = float(reduced[rows, cols].sum()) + float(
sum(lam.sum() for lam in lambdas)
)
best_lower = max(best_lower, lower)
# Feasible recovery => valid upper bound.
assignments, upper = _recover_feasible_nd(cost_tensor, rows, cols)
if upper < best_upper:
best_upper = upper
best_assignments = assignments
gap = best_upper - best_lower
if verbose:
print(
f"Iter {iteration + 1}: LB={best_lower:.4f}, UB={best_upper:.4f}, "
f"Gap={gap:.6f}"
)
if gap < tolerance:
if verbose:
print(f"Converged at iteration {iteration + 1}")
break
# Subgradient: how often each relaxed index is used, minus one.
chosen = np.array(np.unravel_index(best_free[rows, cols], free_shape)).reshape(
len(free_dims), -1
)
subgradient = []
for k, d in enumerate(free_dims):
counts = np.bincount(chosen[k], minlength=dims[d]).astype(float)
subgradient.append(1.0 - counts)
norm_sq = float(sum(float(np.sum(g**2)) for g in subgradient))
if norm_sq <= 0.0:
# No constraint violated: the relaxed solution is feasible and
# therefore optimal.
break
# Polyak step toward the current best upper bound.
step = (best_upper - lower) / norm_sq
if step <= 0.0:
step = 1.0 / (iteration + 1)
for k in range(len(free_dims)):
lambdas[k] = lambdas[k] + step * subgradient[k]
if best_assignments is None:
best_assignments = np.empty((0, n_dims), dtype=np.intp)
best_upper = 0.0
gap = best_upper - best_lower
return AssignmentNDResult(
assignments=best_assignments,
cost=best_upper,
converged=bool(gap < tolerance),
n_iterations=iteration + 1,
gap=gap,
)
[docs]
def auction_assignment_nd(
cost_tensor: NDArray[np.float64],
max_iterations: int = 100,
epsilon: float = 0.01,
verbose: bool = False,
) -> AssignmentNDResult:
"""
Auction algorithm for N-dimensional assignment.
Inspired by the classical auction algorithm for 2D assignment,
adapted to higher dimensions. Objects bid for assignments based
on relative costs.
Parameters
----------
cost_tensor : ndarray
Cost tensor of shape (n1, n2, ..., nk).
max_iterations : int, optional
Maximum iterations (default 100).
epsilon : float, optional
Bid increment (default 0.01). Larger epsilon → fewer iterations,
worse solution; smaller epsilon → more iterations, better solution.
verbose : bool, optional
Print iteration info (default False).
Returns
-------
AssignmentNDResult
Assignments, total cost, convergence info, gap estimate.
Examples
--------
>>> import numpy as np
>>> # 4D assignment: sensors x measurements x tracks x hypotheses
>>> np.random.seed(123)
>>> cost = np.random.rand(2, 3, 3, 2) * 10
>>> result = auction_assignment_nd(cost, max_iterations=50, epsilon=0.1)
>>> len(result.assignments) > 0
True
>>> result.n_iterations <= 50
True
Notes
-----
The algorithm maintains a "price" for each index and allows bidding
(price adjustment) to maximize value. Converges to epsilon-optimal
solution in finite iterations.
"""
dims = cost_tensor.shape
n_dims = len(dims)
# Initialize prices (one per dimension per index)
prices = [np.zeros(dim) for dim in dims]
# Track the best feasible assignment (by actual cost) over all iterations.
# Iteration 0 has zero prices, so the plain greedy solution is included.
best_cost = np.inf
best_assignments = np.empty((0, n_dims), dtype=np.intp)
for iteration in range(max_iterations):
# Compute profit: cost - price penalty
profit = cost_tensor.copy()
for d in range(n_dims):
shape = [1] * n_dims
shape[d] = dims[d]
profit = profit - prices[d].reshape(shape)
# Find best assignment at current prices (greedy)
result = greedy_assignment_nd(profit)
if len(result.assignments) == 0:
break
# Evaluate at actual costs and keep the best solution found
actual_cost = float(np.sum(cost_tensor[tuple(result.assignments.T)]))
if actual_cost < best_cost or (
actual_cost == best_cost and len(result.assignments) > len(best_assignments)
):
best_cost = actual_cost
best_assignments = result.assignments
# Update prices: increase price for "in-demand" indices
demands = [np.zeros(dim) for dim in dims]
for assignment in result.assignments:
for d, idx in enumerate(assignment):
demands[d][idx] += 1
for d in range(n_dims):
prices[d] += epsilon * (demands[d] - 1.0)
if verbose and (iteration + 1) % 10 == 0:
print(f"Iter {iteration + 1}: Cost={actual_cost:.4f}")
if not np.isfinite(best_cost):
best_cost = 0.0
return AssignmentNDResult(
assignments=best_assignments,
cost=best_cost,
converged=True,
n_iterations=iteration + 1,
gap=0.0, # Auction algorithm doesn't track gap formally
)
[docs]
def detect_dimension_conflicts(
assignments: NDArray[np.intp],
dims: Tuple[int, ...],
) -> bool:
"""
Check if assignments violate dimension uniqueness.
For valid assignment, each index should appear at most once per dimension.
Parameters
----------
assignments : ndarray
Array of shape (n_assignments, n_dimensions) with assignments.
dims : tuple
Dimensions of the cost tensor.
Returns
-------
has_conflicts : bool
True if any index appears more than once in any dimension.
Examples
--------
>>> import numpy as np
>>> # Valid assignment: no index repeated in any dimension
>>> assignments = np.array([[0, 0], [1, 1]])
>>> detect_dimension_conflicts(assignments, (3, 3))
False
>>> # Invalid: index 0 used twice in first dimension
>>> assignments = np.array([[0, 0], [0, 1]])
>>> detect_dimension_conflicts(assignments, (3, 3))
True
"""
n_dims = len(dims)
for d in range(n_dims):
indices_in_dim = assignments[:, d]
if len(indices_in_dim) != len(np.unique(indices_in_dim)):
return True
return False
[docs]
class SparseCostTensor:
"""
Sparse representation of N-dimensional cost tensor.
For assignment problems where most entries represent invalid
assignments (infinite cost), storing only valid entries reduces
memory by 50% or more and speeds up greedy algorithms.
Attributes
----------
dims : tuple
Shape of the full tensor (n1, n2, ..., nk).
indices : ndarray
Array of shape (n_valid, n_dims) with valid entry indices.
costs : ndarray
Array of shape (n_valid,) with costs for valid entries.
default_cost : float
Cost for entries not explicitly stored (default: inf).
Examples
--------
>>> import numpy as np
>>> # Create sparse tensor for 10x10x10 problem with 50 valid entries
>>> dims = (10, 10, 10)
>>> valid_indices = np.random.randint(0, 10, size=(50, 3))
>>> valid_costs = np.random.rand(50)
>>> sparse = SparseCostTensor(dims, valid_indices, valid_costs)
>>> sparse.n_valid
50
>>> sparse.sparsity # Fraction of valid entries
0.05
>>> # Convert from dense tensor with inf for invalid
>>> dense = np.full((5, 5, 5), np.inf)
>>> dense[0, 0, 0] = 1.0
>>> dense[1, 1, 1] = 2.0
>>> sparse = SparseCostTensor.from_dense(dense)
>>> sparse.n_valid
2
"""
[docs]
def __init__(
self,
dims: Tuple[int, ...],
indices: NDArray[np.intp],
costs: NDArray[np.float64],
default_cost: float = np.inf,
):
"""
Initialize sparse cost tensor.
Parameters
----------
dims : tuple
Shape of the full tensor.
indices : ndarray
Valid entry indices, shape (n_valid, n_dims).
costs : ndarray
Costs for valid entries, shape (n_valid,).
default_cost : float
Cost for invalid (unstored) entries.
"""
self.dims = dims
self.indices = np.asarray(indices, dtype=np.intp)
self.costs = np.asarray(costs, dtype=np.float64)
self.default_cost = default_cost
# Build lookup for O(1) cost retrieval
self._cost_map: dict[Tuple[int, ...], float] = {}
for i in range(len(self.costs)):
key = tuple(self.indices[i])
self._cost_map[key] = self.costs[i]
@property
def n_dims(self) -> int:
"""Number of dimensions."""
return len(self.dims)
@property
def n_valid(self) -> int:
"""Number of valid (finite cost) entries."""
return len(self.costs)
@property
def sparsity(self) -> float:
"""Fraction of tensor that is valid (0 to 1)."""
total_size = int(np.prod(self.dims))
return self.n_valid / total_size if total_size > 0 else 0.0
@property
def memory_savings(self) -> float:
"""Estimated memory savings vs dense representation (0 to 1)."""
dense_size = np.prod(self.dims) * 8 # 8 bytes per float64
sparse_size = self.n_valid * (8 + self.n_dims * 8) # cost + indices
return max(0, 1 - sparse_size / dense_size) if dense_size > 0 else 0.0
[docs]
def get_cost(self, index: Tuple[int, ...]) -> float:
"""Get cost for a specific index tuple."""
return self._cost_map.get(index, self.default_cost)
[docs]
def to_dense(self) -> NDArray[np.float64]:
"""
Convert to dense tensor representation.
Returns
-------
dense : ndarray
Full tensor with default_cost for unstored entries.
Notes
-----
May use significant memory for large tensors.
"""
dense = np.full(self.dims, self.default_cost, dtype=np.float64)
for i in range(len(self.costs)):
dense[tuple(self.indices[i])] = self.costs[i]
return dense
[docs]
@classmethod
def from_dense(
cls,
dense: NDArray[np.float64],
threshold: float = 1e10,
) -> "SparseCostTensor":
"""
Create sparse tensor from dense array.
Parameters
----------
dense : ndarray
Dense cost tensor.
threshold : float
Entries above this value are considered invalid.
Default 1e10 (catches np.inf and large values).
Returns
-------
SparseCostTensor
Sparse representation.
Examples
--------
>>> import numpy as np
>>> dense = np.array([[[1, np.inf], [np.inf, 2]],
... [[np.inf, 3], [4, np.inf]]])
>>> sparse = SparseCostTensor.from_dense(dense)
>>> sparse.n_valid
4
"""
valid_mask = dense < threshold
indices = np.array(np.where(valid_mask)).T
costs = dense[valid_mask]
return cls(dense.shape, indices, costs, default_cost=np.inf)
[docs]
def greedy_assignment_nd_sparse(
sparse_cost: SparseCostTensor,
max_assignments: Optional[int] = None,
) -> AssignmentNDResult:
"""
Greedy solver for sparse N-dimensional assignment.
Selects minimum-cost tuples from valid entries only, which is much
faster than dense greedy when sparsity < 0.5.
Parameters
----------
sparse_cost : SparseCostTensor
Sparse cost tensor with valid entries only.
max_assignments : int, optional
Maximum number of assignments (default: min(dimensions)).
Returns
-------
AssignmentNDResult
Assignments, total cost, and algorithm info.
Examples
--------
>>> import numpy as np
>>> # Create sparse problem
>>> dims = (10, 10, 10)
>>> # Only 20 valid assignments out of 1000
>>> indices = np.array([[i, i, i] for i in range(10)] +
... [[i, (i+1)%10, (i+2)%10] for i in range(10)])
>>> costs = np.random.rand(20)
>>> sparse = SparseCostTensor(dims, indices, costs)
>>> result = greedy_assignment_nd_sparse(sparse)
>>> result.converged
True
Notes
-----
Time complexity is O(n_valid * log(n_valid)) vs O(total_size * log(total_size))
for dense greedy. For a 10x10x10 tensor with 50 valid entries, this is
50*log(50) vs 1000*log(1000), about 20x faster.
"""
dims = sparse_cost.dims
n_dims = sparse_cost.n_dims
if max_assignments is None:
max_assignments = min(dims)
# Sort valid entries by cost
sorted_indices = np.argsort(sparse_cost.costs)
assignments: List[Tuple[int, ...]] = []
used_indices: List[set[int]] = [set() for _ in range(n_dims)]
total_cost = 0.0
for sorted_idx in sorted_indices:
if len(assignments) >= max_assignments:
break
multi_idx = tuple(sparse_cost.indices[sorted_idx])
# Check if any dimension index is already used
conflict = False
for d, idx in enumerate(multi_idx):
if idx in used_indices[d]:
conflict = True
break
if not conflict:
assignments.append(multi_idx)
total_cost += sparse_cost.costs[sorted_idx]
for d, idx in enumerate(multi_idx):
used_indices[d].add(idx)
assignments_array = np.array(assignments, dtype=np.intp)
if assignments_array.size == 0:
assignments_array = np.empty((0, n_dims), dtype=np.intp)
return AssignmentNDResult(
assignments=assignments_array,
cost=total_cost,
converged=True,
n_iterations=1,
gap=0.0,
)
[docs]
def assignment_nd(
cost: Union[NDArray[np.float64], SparseCostTensor],
method: str = "auto",
max_assignments: Optional[int] = None,
max_iterations: int = 100,
tolerance: float = 1e-6,
epsilon: float = 0.01,
verbose: bool = False,
) -> AssignmentNDResult:
"""
Unified interface for N-dimensional assignment.
Automatically selects between dense and sparse algorithms based on
input type and sparsity.
Parameters
----------
cost : ndarray or SparseCostTensor
Cost tensor (dense) or sparse cost representation.
method : str
Algorithm to use: 'auto', 'greedy', 'relaxation', 'auction'.
'auto' selects greedy for sparse, relaxation for dense.
max_assignments : int, optional
Maximum number of assignments for greedy methods.
max_iterations : int
Maximum iterations for iterative methods.
tolerance : float
Convergence tolerance for relaxation.
epsilon : float
Price increment for auction algorithm.
verbose : bool
Print progress information.
Returns
-------
AssignmentNDResult
Assignment solution.
Examples
--------
>>> import numpy as np
>>> # Dense usage
>>> cost = np.random.rand(4, 4, 4)
>>> result = assignment_nd(cost, method='greedy')
>>> result.converged
True
>>> # Sparse usage (more efficient for large sparse problems)
>>> dense = np.full((20, 20, 20), np.inf)
>>> for i in range(20):
... dense[i, i, i] = np.random.rand()
>>> sparse = SparseCostTensor.from_dense(dense)
>>> result = assignment_nd(sparse, method='auto')
>>> result.converged
True
See Also
--------
greedy_assignment_nd : Dense greedy algorithm.
greedy_assignment_nd_sparse : Sparse greedy algorithm.
relaxation_assignment_nd : Lagrangian relaxation.
auction_assignment_nd : Auction algorithm.
"""
if isinstance(cost, SparseCostTensor):
# Sparse input - use sparse algorithm
if method in ("auto", "greedy"):
return greedy_assignment_nd_sparse(cost, max_assignments)
else:
# Convert to dense for other methods
dense = cost.to_dense()
if method == "relaxation":
return relaxation_assignment_nd(
dense, max_iterations, tolerance, verbose
)
elif method == "auction":
return auction_assignment_nd(
dense, max_iterations, epsilon=epsilon, verbose=verbose
)
else:
raise ValueError(f"Unknown method: {method}")
else:
# Dense input
cost = np.asarray(cost, dtype=np.float64)
if method == "auto":
# Use relaxation for better solutions on dense
return relaxation_assignment_nd(cost, max_iterations, tolerance, verbose)
elif method == "greedy":
return greedy_assignment_nd(cost, max_assignments)
elif method == "relaxation":
return relaxation_assignment_nd(cost, max_iterations, tolerance, verbose)
elif method == "auction":
return auction_assignment_nd(
cost, max_iterations, epsilon=epsilon, verbose=verbose
)
else:
raise ValueError(f"Unknown method: {method}")
__all__ = [
"AssignmentNDResult",
"SparseCostTensor",
"validate_cost_tensor",
"greedy_assignment_nd",
"greedy_assignment_nd_sparse",
"relaxation_assignment_nd",
"auction_assignment_nd",
"detect_dimension_conflicts",
"assignment_nd",
]