113 lines
4.0 KiB
Rust
113 lines
4.0 KiB
Rust
use std::collections::{VecDeque, HashSet};
|
|
use std::iter;
|
|
use std::hash::Hash;
|
|
|
|
use crate::unwrap_or;
|
|
use crate::utils::BoxedIter;
|
|
|
|
/// Two-stage breadth-first search;
|
|
/// Instead of enumerating neighbors before returning a node, it puts visited but not yet
|
|
/// enumerated nodes in a separate queue and only enumerates them to refill the queue of children
|
|
/// one by one once it's empty. This method is preferable for generated graphs because it doesn't
|
|
/// allocate memory for the children until necessary, but it's also probably a bit slower since
|
|
/// it involves additional processing.
|
|
///
|
|
/// # Performance
|
|
/// `T` is cloned twice for each returned value.
|
|
pub fn bfs<T, F, I>(init: T, neighbors: F)
|
|
-> impl Iterator<Item = T>
|
|
where T: Eq + Hash + Clone + std::fmt::Debug,
|
|
F: Fn(T) -> I, I: Iterator<Item = T>
|
|
{
|
|
let mut visited: HashSet<T> = HashSet::new();
|
|
let mut visit_queue: VecDeque<T> = VecDeque::from([init]);
|
|
let mut unpack_queue: VecDeque<T> = VecDeque::new();
|
|
iter::from_fn(move || {
|
|
let next = {loop {
|
|
let next = unwrap_or!(visit_queue.pop_front(); break None);
|
|
if !visited.contains(&next) { break Some(next) }
|
|
}}.or_else(|| loop {
|
|
let unpacked = unwrap_or!(unpack_queue.pop_front(); break None);
|
|
let mut nbv = neighbors(unpacked).filter(|t| !visited.contains(t));
|
|
if let Some(next) = nbv.next() {
|
|
visit_queue.extend(nbv);
|
|
break Some(next)
|
|
}
|
|
})?;
|
|
visited.insert(next.clone());
|
|
unpack_queue.push_back(next.clone());
|
|
Some(next)
|
|
})
|
|
}
|
|
|
|
/// Same as [bfs] but with a recursion depth limit
|
|
///
|
|
/// The main intent is to effectively walk infinite graphs of unknown breadth without making the
|
|
/// recursion depth dependent on the number of nodes. If predictable runtime is more important
|
|
/// than predictable depth, [bfs] with [std::iter::Iterator::take] should be used instead
|
|
pub fn bfs_upto<'a, T: 'a, F: 'a, I: 'a>(init: T, neighbors: F, limit: usize)
|
|
-> impl Iterator<Item = T> + 'a
|
|
where T: Eq + Hash + Clone + std::fmt::Debug,
|
|
F: Fn(T) -> I, I: Iterator<Item = T>
|
|
{
|
|
/// Newtype to store the recursion depth but exclude it from equality comparisons
|
|
/// Because BFS visits nodes in increasing distance order, when a node is visited for the
|
|
/// second time it will never override the earlier version of itself. This is not the case
|
|
/// with Djikstra's algorithm, which can be conceptualised as a "weighted BFS".
|
|
#[derive(Eq, Clone, Debug)]
|
|
struct Wrap<U>(usize, U);
|
|
impl<U: PartialEq> PartialEq for Wrap<U> {
|
|
fn eq(&self, other: &Self) -> bool { self.1.eq(&other.1) }
|
|
}
|
|
impl<U: Hash> Hash for Wrap<U> {
|
|
fn hash<H: std::hash::Hasher>(&self, state: &mut H) { self.1.hash(state) }
|
|
}
|
|
bfs(Wrap(0, init), move |Wrap(dist, t)| -> BoxedIter<Wrap<T>> { // boxed because we branch
|
|
if dist == limit {Box::new(iter::empty())}
|
|
else {Box::new(neighbors(t).map(move |t| Wrap(dist + 1, t)))}
|
|
}).map(|Wrap(_, t)| t)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use itertools::Itertools;
|
|
|
|
use super::*;
|
|
|
|
type Graph = Vec<Vec<usize>>;
|
|
fn neighbors(graph: &Graph, pt: usize) -> impl Iterator<Item = usize> + '_ {
|
|
graph[pt].iter().copied()
|
|
}
|
|
fn from_neighborhood_matrix(matrix: Vec<Vec<usize>>) -> Graph {
|
|
matrix.into_iter().map(|v| {
|
|
v.into_iter().enumerate().filter_map(|(i, ent)| {
|
|
if ent > 1 {panic!("Neighborhood matrices must contain binary values")}
|
|
else if ent == 1 {Some(i)}
|
|
else {None}
|
|
}).collect()
|
|
}).collect()
|
|
}
|
|
|
|
#[test]
|
|
fn test_square() {
|
|
let simple_graph = from_neighborhood_matrix(vec![
|
|
vec![0,1,0,1,1,0,0,0],
|
|
vec![1,0,1,0,0,1,0,0],
|
|
vec![0,1,0,1,0,0,1,0],
|
|
vec![1,0,1,0,0,0,0,1],
|
|
vec![1,0,0,0,0,1,0,1],
|
|
vec![0,1,0,0,1,0,1,0],
|
|
vec![0,0,1,0,0,1,0,1],
|
|
vec![0,0,0,1,1,0,1,0],
|
|
]);
|
|
let scan = bfs(0, |n| neighbors(&simple_graph, n)).collect_vec();
|
|
assert_eq!(scan, vec![0, 1, 3, 4, 2, 5, 7, 6])
|
|
}
|
|
#[test]
|
|
fn test_stringbuilder() {
|
|
let scan = bfs("".to_string(), |s| {
|
|
vec![s.clone()+";", s.clone()+"a", s+"aaa"].into_iter()
|
|
}).take(30).collect_vec();
|
|
println!("{scan:?}")
|
|
}
|
|
} |