Skip to main content

converge_optimization/assignment/
mod.rs

1//! Linear Assignment Problem solvers
2//!
3//! The assignment problem finds the optimal one-to-one matching between
4//! agents and tasks that minimizes total cost.
5//!
6//! ## Problem Definition
7//!
8//! Given:
9//! - n agents and m tasks (typically n = m)
10//! - Cost matrix C where C[i][j] is cost of assigning agent i to task j
11//!
12//! Find:
13//! - Assignment minimizing sum of costs
14//! - Each agent assigned to exactly one task
15//! - Each task assigned to at most one agent
16//!
17//! ## Algorithms
18//!
19//! - [`hungarian`] - O(n⁴) Hungarian algorithm, simple and correct
20//! - [`auction`] - O(n³) Auction algorithm, faster for sparse problems
21//!
22//! ## Example
23//!
24//! ```rust
25//! use converge_optimization::assignment::{solve, AssignmentProblem};
26//!
27//! let problem = AssignmentProblem::from_costs(vec![
28//!     vec![10, 5, 13],
29//!     vec![3, 9, 18],
30//!     vec![14, 8, 7],
31//! ]);
32//!
33//! let solution = solve(&problem).unwrap();
34//! println!("Total cost: {}", solution.total_cost);
35//! for (agent, task) in solution.assignments.iter().enumerate() {
36//!     println!("Agent {} -> Task {}", agent, task);
37//! }
38//! ```
39
40pub mod auction;
41pub mod hungarian;
42
43use crate::{Cost, Error, Result, SolverParams, SolverStats, SolverStatus};
44use serde::{Deserialize, Serialize};
45
46/// An assignment problem instance
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct AssignmentProblem {
49    /// Cost matrix: costs[agent][task]
50    pub costs: Vec<Vec<Cost>>,
51    /// Number of agents
52    pub num_agents: usize,
53    /// Number of tasks
54    pub num_tasks: usize,
55}
56
57impl AssignmentProblem {
58    /// Create a problem from a cost matrix
59    pub fn from_costs(costs: Vec<Vec<Cost>>) -> Self {
60        let num_agents = costs.len();
61        let num_tasks = costs.first().map_or(0, Vec::len);
62        Self {
63            costs,
64            num_agents,
65            num_tasks,
66        }
67    }
68
69    /// Create a square problem from a flat cost vector
70    pub fn from_flat(costs: Vec<Cost>, n: usize) -> Result<Self> {
71        if costs.len() != n * n {
72            return Err(Error::dimension_mismatch(n * n, costs.len()));
73        }
74        let matrix: Vec<Vec<Cost>> = costs.chunks(n).map(|c| c.to_vec()).collect();
75        Ok(Self::from_costs(matrix))
76    }
77
78    /// Check if the problem is square (n agents = n tasks)
79    pub fn is_square(&self) -> bool {
80        self.num_agents == self.num_tasks
81    }
82
83    /// Get cost of assigning agent to task
84    pub fn cost(&self, agent: usize, task: usize) -> Cost {
85        self.costs[agent][task]
86    }
87
88    /// Validate the problem
89    pub fn validate(&self) -> Result<()> {
90        if self.num_agents == 0 {
91            return Err(Error::invalid_input("no agents"));
92        }
93        if self.num_tasks == 0 {
94            return Err(Error::invalid_input("no tasks"));
95        }
96        for row in &self.costs {
97            if row.len() != self.num_tasks {
98                return Err(Error::dimension_mismatch(self.num_tasks, row.len()));
99            }
100        }
101        Ok(())
102    }
103}
104
105/// Solution to an assignment problem
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct AssignmentSolution {
108    /// Assignment: assignments[agent] = task
109    pub assignments: Vec<usize>,
110    /// Total cost of the assignment
111    pub total_cost: Cost,
112    /// Solver status
113    pub status: SolverStatus,
114    /// Solver statistics
115    pub stats: SolverStats,
116}
117
118impl AssignmentSolution {
119    /// Get the task assigned to an agent
120    pub fn task_for_agent(&self, agent: usize) -> Option<usize> {
121        self.assignments.get(agent).copied()
122    }
123
124    /// Iterate over (agent, task) pairs
125    pub fn iter(&self) -> impl Iterator<Item = (usize, usize)> + '_ {
126        self.assignments.iter().enumerate().map(|(a, &t)| (a, t))
127    }
128}
129
130/// Trait for assignment solvers
131pub trait AssignmentSolver {
132    /// Solve the assignment problem
133    fn solve(
134        &self,
135        problem: &AssignmentProblem,
136        params: &SolverParams,
137    ) -> Result<AssignmentSolution>;
138
139    /// Solver name
140    fn name(&self) -> &'static str;
141}
142
143/// Solve an assignment problem using the default solver (Hungarian)
144pub fn solve(problem: &AssignmentProblem) -> Result<AssignmentSolution> {
145    solve_with_params(problem, &SolverParams::default())
146}
147
148/// Solve an assignment problem with custom parameters
149pub fn solve_with_params(
150    problem: &AssignmentProblem,
151    params: &SolverParams,
152) -> Result<AssignmentSolution> {
153    problem.validate()?;
154    hungarian::HungarianSolver.solve(problem, params)
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn test_simple_assignment() {
163        let problem =
164            AssignmentProblem::from_costs(vec![vec![10, 5, 13], vec![3, 9, 18], vec![14, 8, 7]]);
165
166        let solution = solve(&problem).unwrap();
167        assert_eq!(solution.status, SolverStatus::Optimal);
168        // Optimal: agent 0->task 1 (5), agent 1->task 0 (3), agent 2->task 2 (7) = 15
169        assert_eq!(solution.total_cost, 15);
170    }
171
172    #[test]
173    fn test_from_flat() {
174        let costs = vec![1, 2, 3, 4];
175        let problem = AssignmentProblem::from_flat(costs, 2).unwrap();
176        assert_eq!(problem.num_agents, 2);
177        assert_eq!(problem.num_tasks, 2);
178        assert_eq!(problem.cost(0, 0), 1);
179        assert_eq!(problem.cost(1, 1), 4);
180    }
181}