Skip to content

Commit

Permalink
Further optimize DFS/BFS algorithms
Browse files Browse the repository at this point in the history
  • Loading branch information
@regexident
regexident committed Oct 23, 2024
1 parent ca7da87 commit f5a885d
Show file tree
Hide file tree
Showing 2 changed files with 166 additions and 68 deletions.
117 changes: 82 additions & 35 deletions crates/algos/src/bfs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,61 +6,85 @@ use crate::prelude::*;

pub fn bfs_directed<NI, G>(
graph: &G,
node_id: NI,
node_ids: impl IntoIterator<Item = NI>,
direction: Direction,
) -> DirectedBreadthFirst<'_, G, NI>
where
NI: Idx + Hash,
G: Graph<NI> + DirectedDegrees<NI> + DirectedNeighbors<NI> + Sync,
{
DirectedBreadthFirst::new(graph, node_id, direction)
DirectedBreadthFirst::new(graph, node_ids, direction)
}

pub struct DirectedBreadthFirst<'a, G, NI> {
graph: &'a G,
seen: BitVec<usize>,
visited: BitVec<usize>,
queue: VecDeque<NI>,
direction: Direction,
}

impl<'a, G, NI> DirectedBreadthFirst<'a, G, NI>
where
NI: Idx + Hash,
NI: Idx + Hash + std::fmt::Debug,
G: Graph<NI> + DirectedNeighbors<NI> + Sync,
{
pub fn new(graph: &'a G, node_id: NI, direction: Direction) -> Self {
pub fn new(graph: &'a G, node_ids: impl IntoIterator<Item = NI>, direction: Direction) -> Self {
let bitvec = BitVec::repeat(false, graph.node_count().index());
let visited = bitvec.clone();

let mut seen = bitvec;
let mut queue = VecDeque::new();
Self::enqueue_into(&mut seen, &mut queue, node_ids);

Self {
graph,
visited: BitVec::repeat(false, graph.node_count().index()),
queue: VecDeque::from_iter([node_id]),
seen,
visited,
queue,
direction,
}
}

fn dequeue(&mut self) -> Option<NI> {
loop {
let node_id = self.queue.pop_front()?;

if !self.visited.replace(node_id.index(), true) {
return Some(node_id);
}
}
}

fn enqueue_out_neighbors(&mut self, node_id: NI) {
let neighbors = self
fn enqueue_into(
seen: &mut BitVec<usize>,
queue: &mut VecDeque<NI>,
node_ids: impl IntoIterator<Item = NI>,
) {
for node_id in node_ids {
if !seen.replace(node_id.index(), true) {
queue.push_back(node_id);
}
}
}

fn enqueue_out_neighbors_of(&mut self, node_id: NI) {
let node_ids = self
.graph
.out_neighbors(node_id)
.filter(|&node_id| !self.visited[node_id.index()]);
self.queue.extend(neighbors);
.copied()
.filter(|node_id| !self.visited[node_id.index()]);

Self::enqueue_into(&mut self.seen, &mut self.queue, node_ids);
}

fn enqueue_in_neighbors(&mut self, node_id: NI) {
let neighbors = self
fn enqueue_in_neighbors_of(&mut self, node_id: NI) {
let node_ids = self
.graph
.in_neighbors(node_id)
.filter(|&node_id| !self.visited[node_id.index()]);
self.queue.extend(neighbors);
.copied()
.filter(|node_id| !self.visited[node_id.index()]);

Self::enqueue_into(&mut self.seen, &mut self.queue, node_ids);
}
}

Expand All @@ -75,28 +99,32 @@ where
let node_id = self.dequeue()?;

match self.direction {
Direction::Outgoing => self.enqueue_out_neighbors(node_id),
Direction::Incoming => self.enqueue_in_neighbors(node_id),
Direction::Outgoing => self.enqueue_out_neighbors_of(node_id),
Direction::Incoming => self.enqueue_in_neighbors_of(node_id),
Direction::Undirected => {
self.enqueue_out_neighbors(node_id);
self.enqueue_in_neighbors(node_id);
self.enqueue_out_neighbors_of(node_id);
self.enqueue_in_neighbors_of(node_id);
}
}

Some(node_id)
}
}

pub fn bfs_undirected<NI, G>(graph: &G, node_id: NI) -> UndirectedBreadthFirst<'_, G, NI>
pub fn bfs_undirected<NI, G>(
graph: &G,
node_ids: impl IntoIterator<Item = NI>,
) -> UndirectedBreadthFirst<'_, G, NI>
where
NI: Idx + Hash,
G: Graph<NI> + UndirectedDegrees<NI> + UndirectedNeighbors<NI> + Sync,
{
UndirectedBreadthFirst::new(graph, node_id)
UndirectedBreadthFirst::new(graph, node_ids)
}

pub struct UndirectedBreadthFirst<'a, G, NI> {
graph: &'a G,
seen: BitVec<usize>,
visited: BitVec<usize>,
queue: VecDeque<NI>,
}
Expand All @@ -106,11 +134,19 @@ where
NI: Idx + Hash + std::fmt::Debug,
G: Graph<NI> + UndirectedNeighbors<NI> + Sync,
{
pub fn new(graph: &'a G, node_id: NI) -> Self {
pub fn new(graph: &'a G, node_ids: impl IntoIterator<Item = NI>) -> Self {
let bitvec = BitVec::repeat(false, graph.node_count().index());
let visited = bitvec.clone();

let mut seen = bitvec;
let mut queue = VecDeque::new();
Self::enqueue_into(&mut seen, &mut queue, node_ids);

Self {
graph,
visited: BitVec::repeat(false, graph.node_count().index()),
queue: VecDeque::from_iter([node_id]),
seen,
visited,
queue,
}
}

Expand All @@ -124,15 +160,26 @@ where
}
}

fn enqueue_neighbors(&mut self, node_id: NI) {
let neighbors = self
fn enqueue_into(
seen: &mut BitVec<usize>,
queue: &mut VecDeque<NI>,
node_ids: impl IntoIterator<Item = NI>,
) {
for node_id in node_ids {
if !seen.replace(node_id.index(), true) {
queue.push_back(node_id);
}
}
}

fn enqueue_neighbors_of(&mut self, node_id: NI) {
let node_ids = self
.graph
.neighbors(node_id)
.copied()
.filter(|&node_id| !self.visited[node_id.index()]);

let neighbors: Vec<_> = neighbors.collect();

self.queue.extend(neighbors);
Self::enqueue_into(&mut self.seen, &mut self.queue, node_ids);
}
}

Expand All @@ -146,18 +193,18 @@ where
fn next(&mut self) -> Option<Self::Item> {
let node_id = self.dequeue()?;

self.enqueue_neighbors(node_id);
self.enqueue_neighbors_of(node_id);

Some(node_id)
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::prelude::{CsrLayout, GraphBuilder};
use graph::prelude::{CsrLayout, GraphBuilder};

use super::*;

mod directed {
use super::*;

Expand All @@ -168,7 +215,7 @@ mod tests {
.edges(vec![(0, 1), (0, 2), (1, 2), (1, 3), (2, 3), (2, 1), (3, 1)])
.build();

let actual: Vec<usize> = bfs_directed(&graph, 0, Direction::Outgoing).collect();
let actual: Vec<usize> = bfs_directed(&graph, [0], Direction::Outgoing).collect();
let expected: Vec<usize> = vec![0, 1, 2, 3];

assert_eq!(actual, expected);
Expand All @@ -181,7 +228,7 @@ mod tests {
.edges(vec![(0, 1), (0, 2), (1, 2), (1, 3), (2, 1), (2, 1), (3, 1)])
.build();

let actual: Vec<usize> = bfs_directed(&graph, 0, Direction::Outgoing).collect();
let actual: Vec<usize> = bfs_directed(&graph, [0], Direction::Outgoing).collect();
let expected: Vec<usize> = vec![0, 1, 2, 3];

assert_eq!(actual, expected);
Expand All @@ -195,7 +242,7 @@ mod tests {
.edges(vec![(0, 1), (0, 2), (1, 2), (1, 3), (2, 3), (2, 1), (3, 1)])
.build();

let actual: Vec<usize> = bfs_undirected(&graph, 0).collect();
let actual: Vec<usize> = bfs_undirected(&graph, [0]).collect();
let expected: Vec<usize> = vec![0, 1, 2, 3];

assert_eq!(actual, expected);
Expand Down
Loading

0 comments on commit f5a885d

Please sign in to comment.