Merge #309

309: Add reachable() to Matrix r=samueltardieu a=samueltardieu



Co-authored-by: Samuel Tardieu <[email protected]>

Author: bors[bot]

src/matrix.rs

@@ -1,9 +1,11 @@
 //! Matrix of an arbitrary type and utilities to rotate, transpose, etc.
 
+use crate::directed::bfs::bfs_reach;
 use crate::utils::uint_sqrt;
 use itertools::iproduct;
 use itertools::Itertools;
 use num_traits::Signed;
+use std::collections::BTreeSet;
 use std::error::Error;
 use std::fmt;
 use std::ops::{Deref, DerefMut, Index, IndexMut, Neg, Range};
@@ -516,6 +518,27 @@ impl<C> Matrix<C> {
     pub fn values_mut(&mut self) -> IterMut<C> {
         self.data.iter_mut()
     }
+
+    /// Return a set of the indices reachable from a candidate starting point
+    /// and for which the given predicate is valid. This can be used for example
+    /// to implement a flood-filling algorithm. Since the indices are collected
+    /// into a vector, they can later be used without keeping a reference on the
+    /// matrix itself, e.g., to modify the matrix.
+    pub fn reachable<P>(
+        &self,
+        start: (usize, usize),
+        diagonals: bool,
+        mut predicate: P,
+    ) -> BTreeSet<(usize, usize)>
+    where
+        P: FnMut((usize, usize)) -> bool,
+        P: Copy,
+    {
+        bfs_reach(start, move |&n| {
+            self.neighbours(n, diagonals).filter(move |&n| predicate(n))
+        })
+        .collect()
+    }
 }
 
 impl<C> Index<(usize, usize)> for Matrix<C> {

tests/matrix.rs

@@ -379,6 +379,23 @@ fn empty_neighbours() {
     assert_eq!(m.neighbours((10, 10), true).collect::<Vec<_>>(), vec![]);
 }
 
+#[test]
+fn reachable() {
+    let m = matrix![[0, 1, 2], [3, 4, 5], [6, 7, 8]];
+
+    let indices = m.reachable((1, 0), false, |n| m[n] % 4 != 0);
+    assert_eq!(
+        indices.into_iter().collect::<Vec<_>>(),
+        vec![(1, 0), (2, 0), (2, 1)]
+    );
+
+    let indices = m.reachable((1, 0), true, |n| m[n] % 4 != 0);
+    assert_eq!(
+        indices.into_iter().collect::<Vec<_>>(),
+        vec![(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]
+    );
+}
+
 #[test]
 fn from_rows() {
     let m = Matrix::from_rows((1..3).map(|n| (1..5).map(move |x| x * n))).unwrap();