1  use std::{cmp::min, iter::FusedIterator, num::NonZeroUsize};
  2  
  3  pub trait NumDigits {
  4      fn num_digits(&self, radix: u32) -> u8;
  5  }
  6  
  7  impl NumDigits for u64 {
  8      fn num_digits(&self, radix: u32) -> u8 {
  9          (self.ilog(radix as u64) + 1) as u8
 10      }
 11  }
 12  
 13  pub fn iter_neighbors<'a, Outer, Inner, T: 'a>(
 14      grid: &'a Outer,
 15      row: usize,
 16      col: usize,
 17  ) -> impl Iterator<Item = &'a T>
 18  where
 19      Outer: AsRef<[Inner]>,
 20      Inner: AsRef<[T]> + 'a,
 21  {
 22      let rows = grid.as_ref().len();
 23      let cols = grid.as_ref()[0].as_ref().len();
 24      assert!(row < rows);
 25      assert!(col < cols);
 26  
 27      grid.as_ref()[row.saturating_sub(1)..=min(row + 1, rows - 1)]
 28          .iter()
 29          .enumerate()
 30          .flat_map(move |(r, tiles)| {
 31              tiles.as_ref()[col.saturating_sub(1)..=min(col + 1, cols - 1)]
 32                  .iter()
 33                  .enumerate()
 34                  .filter(move |(c, _)| {
 35                      r + row.saturating_sub(1) != row || c + col.saturating_sub(1) != col
 36                  })
 37                  .map(|(_, tile)| tile)
 38          })
 39  }
 40  
 41  pub fn iter_neighbors_mut<'a, Outer, Inner, T: 'a>(
 42      grid: &'a mut Outer,
 43      row: usize,
 44      col: usize,
 45  ) -> impl Iterator<Item = &'a mut T>
 46  where
 47      Outer: AsMut<[Inner]>,
 48      Inner: AsMut<[T]> + 'a,
 49  {
 50      let rows = grid.as_mut().len();
 51      let cols = grid.as_mut()[0].as_mut().len();
 52      assert!(row < rows);
 53      assert!(col < cols);
 54  
 55      grid.as_mut()[row.saturating_sub(1)..=min(row + 1, rows - 1)]
 56          .iter_mut()
 57          .enumerate()
 58          .flat_map(move |(r, tiles)| {
 59              tiles.as_mut()[col.saturating_sub(1)..=min(col + 1, cols - 1)]
 60                  .iter_mut()
 61                  .enumerate()
 62                  .filter(move |(c, _)| {
 63                      r + row.saturating_sub(1) != row || c + col.saturating_sub(1) != col
 64                  })
 65                  .map(|(_, tile)| tile)
 66          })
 67  }
 68  
 69  /// iterator over the coordinates adjacent to a given row and column in a given grid.
 70  ///
 71  /// this doesn't retain a reference to the grid on which it operates, so this is useful for indexed
 72  /// mutations. it is also [`ExactSizeIterator`].
 73  #[derive(Debug, Clone)]
 74  pub struct NeighborCoords {
 75      // TODO space optimize this with i8 offsets
 76      init_row: usize,
 77      current_row: usize,
 78      row_bound: usize,
 79      init_col: usize,
 80      current_col: usize,
 81      col_bound: usize,
 82  }
 83  
 84  impl ExactSizeIterator for NeighborCoords {}
 85  impl FusedIterator for NeighborCoords {}
 86  
 87  impl NeighborCoords {
 88      pub fn new(init_row: usize, init_col: usize, row_bound: usize, col_bound: usize) -> Self {
 89          NeighborCoords {
 90              init_row,
 91              current_row: init_row.saturating_sub(1),
 92              row_bound,
 93              init_col,
 94              current_col: init_col.saturating_sub(1),
 95              col_bound,
 96          }
 97      }
 98  
 99      // TODO genericize this :3
100      pub fn in_grid<_T>(row: usize, col: usize, grid: &[Vec<_T>]) -> Self {
101          assert!(row < grid.len());
102          assert!(col < grid[0].len());
103          NeighborCoords::new(row, col, grid.len(), grid[0].len())
104      }
105  
106      fn row_lower_bound(&self) -> usize {
107          self.init_row.saturating_sub(1)
108      }
109  
110      fn col_lower_bound(&self) -> usize {
111          self.init_col.saturating_sub(1)
112      }
113  
114      fn row_upper_bound(&self) -> usize {
115          min(self.row_bound, self.init_row + 2)
116      }
117  
118      fn col_upper_bound(&self) -> usize {
119          min(self.col_bound, self.init_col + 2)
120      }
121  
122      fn rows(&self) -> usize {
123          self.row_upper_bound() - self.row_lower_bound()
124      }
125  
126      fn cols(&self) -> usize {
127          self.col_upper_bound() - self.col_lower_bound()
128      }
129  }
130  
131  impl Iterator for NeighborCoords {
132      type Item = (usize, usize);
133  
134      fn next(&mut self) -> Option<Self::Item> {
135          if self.current_row == self.row_upper_bound() {
136              None
137          } else if self.current_col == self.col_upper_bound() {
138              self.current_col = self.col_lower_bound();
139              self.current_row += 1;
140              self.next()
141          } else if self.current_row == self.init_row && self.current_col == self.init_col {
142              self.current_col += 1;
143              self.next()
144          } else {
145              self.current_col += 1;
146              Some((self.current_row, self.current_col - 1))
147          }
148      }
149  
150      fn size_hint(&self) -> (usize, Option<usize>) {
151          let total = self.rows() * self.cols();
152          let visited = (self.current_row - self.row_lower_bound()) * self.cols()
153              + (self.current_col - self.col_lower_bound());
154  
155          // subtract 1 if we haven't skipped the middle yet
156          if (self.init_row - self.row_lower_bound()) * self.cols()
157              + (self.init_col - self.col_lower_bound())
158              >= (self.current_row - self.row_lower_bound()) * self.cols()
159                  + (self.current_col - self.col_lower_bound())
160          {
161              (total - visited - 1, Some(total - visited - 1))
162          } else {
163              (total - visited, Some(total - visited))
164          }
165      }
166  }
167  
168  #[derive(Debug, Clone, Copy, Default)]
169  struct DisjoinSetsNode {
170      parent_plus_one: Option<NonZeroUsize>,
171      rank: u32,
172  }
173  
174  // TODO maybe reimplement with refcell so api isnt all mutating functions
175  #[derive(Debug, Clone)]
176  pub struct DisjointSets {
177      nodes: Vec<DisjoinSetsNode>,
178  }
179  
180  impl DisjointSets {
181      pub fn new(size: usize) -> Self {
182          DisjointSets {
183              nodes: vec![Default::default(); size],
184          }
185      }
186  
187      /// returns the set representative of `key` if `key` is a member of a set.
188      pub fn find(&mut self, key: usize) -> usize {
189          let mut root = key;
190          while let Some(parent_plus_one) = self.nodes[root].parent_plus_one {
191              root = parent_plus_one.get() - 1;
192          }
193  
194          let mut y = key;
195          while y != root {
196              // this is suspect
197              let temp = self.nodes[y];
198              self.nodes[y] = DisjoinSetsNode {
199                  parent_plus_one: Some(NonZeroUsize::new(root + 1).unwrap()),
200                  rank: temp.rank,
201              };
202              y = temp.parent_plus_one.unwrap().get() - 1;
203          }
204  
205          root
206      }
207  
208      /// combines the sets that contain `k1` and `k2`.
209      pub fn union(&mut self, k1: usize, k2: usize) -> bool {
210          let r1 = self.find(k1);
211          let r2 = self.find(k2);
212  
213          if r1 == r2 {
214              return false;
215          }
216  
217          if self.nodes[r1].rank < self.nodes[r2].rank {
218              self.nodes[r1].parent_plus_one = Some(NonZeroUsize::new(r2 + 1).unwrap());
219              if self.nodes[r2].rank == self.nodes[r2].rank {
220                  self.nodes[r1].rank += 1;
221              }
222          } else {
223              self.nodes[r2].parent_plus_one = Some(NonZeroUsize::new(r1 + 1).unwrap());
224              if self.nodes[r1].rank == self.nodes[r2].rank {
225                  self.nodes[r2].rank += 1;
226              }
227          }
228  
229          true
230      }
231  
232      pub fn collect(&mut self) -> Vec<Vec<usize>> {
233          let mut sets = vec![vec![]; self.nodes.len()];
234  
235          for i in 0..self.nodes.len() {
236              sets[self.find(i)].push(i);
237          }
238  
239          sets.into_iter().filter(|set| !set.is_empty()).collect()
240      }
241  }
242  
243  #[cfg(test)]
244  mod tests {
245      use super::*;
246  
247      #[test]
248      fn test_yields_expected() {
249          let grid = vec![vec![(), (), ()], vec![(), (), ()], vec![(), (), ()]];
250  
251          let actual = NeighborCoords::in_grid(1, 0, &grid).collect::<Vec<_>>();
252          let expected = vec![(0, 0), (0, 1), (1, 1), (2, 0), (2, 1)];
253          assert_eq!(actual, expected);
254  
255          let actual = NeighborCoords::in_grid(0, 1, &grid).collect::<Vec<_>>();
256          let expected = vec![(0, 0), (0, 2), (1, 0), (1, 1), (1, 2)];
257          assert_eq!(actual, expected);
258  
259          let actual = NeighborCoords::in_grid(1, 1, &grid).collect::<Vec<_>>();
260          let expected = vec![
261              (0, 0),
262              (0, 1),
263              (0, 2),
264              (1, 0),
265              (1, 2),
266              (2, 0),
267              (2, 1),
268              (2, 2),
269          ];
270          assert_eq!(actual, expected);
271  
272          let actual = NeighborCoords::in_grid(1, 2, &grid).collect::<Vec<_>>();
273          let expected = vec![(0, 1), (0, 2), (1, 1), (2, 1), (2, 2)];
274          assert_eq!(actual, expected);
275  
276          let actual = NeighborCoords::in_grid(2, 1, &grid).collect::<Vec<_>>();
277          let expected = vec![(1, 0), (1, 1), (1, 2), (2, 0), (2, 2)];
278          assert_eq!(actual, expected);
279  
280          let big_grid = vec![
281              vec![0, 0, 0, 0, 0, 0],
282              vec![0, 0, 0, 0, 0, 0],
283              vec![0, 0, 0, 0, 0, 0],
284              vec![0, 0, 0, 0, 0, 0],
285              vec![0, 0, 1, 0, 0, 0],
286              vec![0, 0, 0, 0, 0, 0],
287          ];
288  
289          let actual = NeighborCoords::in_grid(4, 2, &big_grid).collect::<Vec<_>>();
290          let expected = vec![
291              (3, 1),
292              (3, 2),
293              (3, 3),
294              (4, 1),
295              (4, 3),
296              (5, 1),
297              (5, 2),
298              (5, 3),
299          ];
300          assert_eq!(actual, expected);
301      }
302  
303      #[test]
304      fn test_exact_size() {
305          let grid = vec![
306              vec![0, 0, 0, 0, 0, 0],
307              vec![0, 0, 0, 0, 0, 0],
308              vec![0, 0, 0, 0, 0, 0],
309              vec![0, 0, 0, 0, 0, 0],
310              vec![0, 0, 0, 0, 0, 0],
311              vec![0, 0, 0, 0, 0, 0],
312          ];
313  
314          let mut iter = NeighborCoords::in_grid(1, 0, &grid);
315          assert_eq!(iter.len(), 5);
316          iter.next();
317          assert_eq!(iter.len(), 4);
318          iter.next();
319          assert_eq!(iter.len(), 3);
320          iter.next();
321          assert_eq!(iter.len(), 2);
322          iter.next();
323          assert_eq!(iter.len(), 1);
324          iter.next();
325          assert_eq!(iter.len(), 0);
326          assert_eq!(iter.next(), None);
327  
328          let mut iter = NeighborCoords::in_grid(0, 1, &grid);
329          assert_eq!(iter.len(), 5);
330          iter.next();
331          assert_eq!(iter.len(), 4);
332          iter.next();
333          assert_eq!(iter.len(), 3);
334          iter.next();
335          assert_eq!(iter.len(), 2);
336          iter.next();
337          assert_eq!(iter.len(), 1);
338          iter.next();
339          assert_eq!(iter.len(), 0);
340          assert_eq!(iter.next(), None);
341  
342          let mut iter = NeighborCoords::in_grid(4, 5, &grid);
343          assert_eq!(iter.len(), 5);
344          iter.next();
345          assert_eq!(iter.len(), 4);
346          iter.next();
347          assert_eq!(iter.len(), 3);
348          iter.next();
349          assert_eq!(iter.len(), 2);
350          iter.next();
351          assert_eq!(iter.len(), 1);
352          iter.next();
353          assert_eq!(iter.len(), 0);
354          assert_eq!(iter.next(), None);
355  
356          let mut iter = NeighborCoords::in_grid(5, 4, &grid);
357          assert_eq!(iter.len(), 5);
358          iter.next();
359          assert_eq!(iter.len(), 4);
360          iter.next();
361          assert_eq!(iter.len(), 3);
362          iter.next();
363          assert_eq!(iter.len(), 2);
364          iter.next();
365          assert_eq!(iter.len(), 1);
366          iter.next();
367          assert_eq!(iter.len(), 0);
368          assert_eq!(iter.next(), None);
369  
370          let mut iter = NeighborCoords::in_grid(4, 2, &grid);
371          assert_eq!(iter.len(), 8);
372          iter.next();
373          assert_eq!(iter.len(), 7);
374          iter.next();
375          assert_eq!(iter.len(), 6);
376          iter.next();
377          assert_eq!(iter.len(), 5);
378          iter.next();
379          assert_eq!(iter.len(), 4);
380          iter.next();
381          assert_eq!(iter.len(), 3);
382          iter.next();
383          assert_eq!(iter.len(), 2);
384          iter.next();
385          assert_eq!(iter.len(), 1);
386          iter.next();
387          assert_eq!(iter.len(), 0);
388          assert_eq!(iter.next(), None);
389      }
390  }