Skip to main content

munkres/
mark_matrix.rs

1use crate::{Position, SquareMatrix};
2
3#[derive(Debug)]
4pub struct MarkMatrix {
5    marks: SquareMatrix<Mark>,
6}
7
8#[derive(Clone, Copy, PartialEq, Eq, Debug)]
9#[repr(u8)]
10enum Mark {
11    None,
12    Star,
13    Prime,
14}
15
16impl MarkMatrix {
17    pub fn new(n: usize) -> Self {
18        Self {
19            marks: SquareMatrix::from_shape_fn((n, n), |_| Mark::None),
20        }
21    }
22
23    #[inline]
24    pub fn n(&self) -> usize {
25        self.marks.shape()[0]
26    }
27
28    #[inline]
29    fn get_mark(&self, pos: Position) -> Mark {
30        self.marks[(pos.row, pos.column)]
31    }
32
33    #[inline]
34    fn set_mark(&mut self, pos: Position, mark: Mark) {
35        self.marks[(pos.row, pos.column)] = mark;
36    }
37
38    pub fn toggle_star(&mut self, pos: Position) {
39        let cell = self.marks.get_mut((pos.row, pos.column)).unwrap();
40        if *cell == Mark::Star {
41            *cell = Mark::None;
42        } else {
43            *cell = Mark::Star;
44        }
45    }
46
47    pub fn star(&mut self, pos: Position) {
48        self.set_mark(pos, Mark::Star);
49    }
50
51    pub fn prime(&mut self, pos: Position) {
52        self.set_mark(pos, Mark::Prime);
53    }
54
55    pub fn is_star(&self, pos: Position) -> bool {
56        match self.get_mark(pos) {
57            Mark::Star => true,
58            _ => false,
59        }
60    }
61
62    pub fn is_prime(&self, pos: Position) -> bool {
63        match self.get_mark(pos) {
64            Mark::Prime => true,
65            _ => false,
66        }
67    }
68
69    #[cfg(test)]
70    pub fn is_none(&self, pos: Position) -> bool {
71        match self.get_mark(pos) {
72            Mark::None => true,
73            _ => false,
74        }
75    }
76
77    #[inline]
78    pub fn each_star<F>(&self, mut f: F)
79    where
80        F: FnMut(Position),
81    {
82        for (row, row_data) in self.marks.genrows().into_iter().enumerate() {
83            for (column, &cell) in row_data.iter().enumerate() {
84                if cell == Mark::Star {
85                    f(Position { row, column });
86                }
87            }
88        }
89    }
90
91    fn find_first_mark_in_row(&self, row: usize, mark: Mark) -> Option<usize> {
92        self.marks
93            .row(row)
94            .iter()
95            .enumerate()
96            .find(|(_, &cell)| cell == mark)
97            .map(|(column, _)| column)
98    }
99
100    fn find_first_mark_in_column(&self, column: usize, mark: Mark) -> Option<usize> {
101        self.marks
102            .column(column)
103            .iter()
104            .enumerate()
105            .find(|(_, &cell)| cell == mark)
106            .map(|(row, _)| row)
107    }
108
109    pub fn find_first_star_in_row(&self, row: usize) -> Option<usize> {
110        self.find_first_mark_in_row(row, Mark::Star)
111    }
112
113    pub fn find_first_prime_in_row(&self, row: usize) -> Option<usize> {
114        self.find_first_mark_in_row(row, Mark::Prime)
115    }
116
117    pub fn find_first_star_in_column(&self, column: usize) -> Option<usize> {
118        self.find_first_mark_in_column(column, Mark::Star)
119    }
120
121    pub fn clear_primes(&mut self) {
122        for cell in self.marks.iter_mut() {
123            if *cell == Mark::Prime {
124                *cell = Mark::None
125            }
126        }
127    }
128}