use std::{cmp::Ordering, collections::VecDeque};
use crate::union_find::UnionFind;
pub type Distance = usize;
pub trait Graph {
type Node;
type Edge: Edge;
/// All nodes.
fn nodes(&self) -> impl Iterator<Item = &Self::Node>;
/// All edges.
fn edges(&self) -> impl Iterator<Item = &Self::Edge>;
/// Indices of neighbors reachable from node.
fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>>;
/// Indices of neighbors that can reach node.
/// For directed graphs, returns same result as neighbor_idxs
fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>>;
/// All edges going to or from node.
/// For directed graphs, mind the from and to distinction.
fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>>;
/// Look if there is an edge between nodes from and to.
/// It is expected this function will be overriden with more performant version.
fn has_edge(&self, from: usize, to: usize) -> bool {
self.edges()
.any(|edge| edge.from_node() == from && edge.to_node() == to)
}
/// Find an edge that connects nodes from and to, if it exists.
/// If it doesn't, return none.
/// It is expected this function will be overriden with more performant version.
fn get_edge_between(&self, from: usize, to: usize) -> Option<usize> {
self.edges_of_idxs(from)
.map(|edges| edges
.filter(|&edge| self.edge(edge).unwrap().to_node() == to)
.next())
.flatten()
}
/// Get a single edge at the given index.
fn edge(&self, id: usize) -> Option<&Self::Edge> {
self.edges().skip(id).next()
}
/// Get a single node at the given index.
fn node(&self, id: usize) -> Option<&Self::Node> {
self.nodes().skip(id).next()
}
}
pub trait MutGraph: Graph {
fn nodes_mut(&mut self) -> impl Iterator<Item = &mut Self::Node>;
fn edges_mut(&mut self) -> impl Iterator<Item = &mut Self::Edge>;
fn node_mut(&mut self, id: usize) -> Option<&mut Self::Node> {
self.nodes_mut().skip(id).next()
}
fn edge_mut(&mut self, id: usize) -> Option<&mut Self::Edge> {
self.edges_mut().skip(id).next()
}
}
pub trait WeightedEdge {
type Cost: PartialOrd;
fn cost(&self) -> Self::Cost;
}
/// An edge.
pub trait Edge {
/// Index of a node this edge goes from.
/// For undirected graphs, from_node and to_node have the same meaning.
fn from_node(&self) -> usize;
/// Index of a node this edge goes to.
/// For undirected graphs, from_node and to_node have the same meaning.
fn to_node(&self) -> usize;
}
/// An edge that might be reversed easily.
pub trait ReversibleEdge: Edge {
fn reverse(self) -> Self;
}
#[derive(Debug, Clone, PartialEq)]
pub struct GenericEdge {
from: usize,
to: usize,
}
impl GenericEdge {
pub fn new(from: usize, to: usize) -> Self {
Self {
from,
to
}
}
}
impl From<(usize, usize)> for GenericEdge {
fn from(value: (usize, usize)) -> Self {
Self {
from: value.0,
to: value.1
}
}
}
impl Edge for GenericEdge {
fn from_node(&self) -> usize {
self.from
}
fn to_node(&self) -> usize {
self.to
}
}
impl ReversibleEdge for GenericEdge {
fn reverse(mut self) -> Self {
(self.from, self.to) = (self.to, self.from);
self
}
}
/// A directed graph that owns nodes and edges of
/// specific types given by generics.
#[derive(Debug, Clone, PartialEq)]
pub struct GenericDirectedGraph<T, TEdge: Edge>
{
nodes: Vec<T>,
edges: Vec<TEdge>,
node_edges: Vec<Vec<usize>>,
reverse_node_edges: Vec<Vec<usize>>,
}
impl<T, TEdge: Edge> GenericDirectedGraph<T, TEdge> {
pub fn new(nodes: Vec<T>) -> Self {
let nodes_len = nodes.len();
Self {
nodes,
edges: vec![],
node_edges: vec![vec![]; nodes_len],
reverse_node_edges: vec![vec![]; nodes_len],
}
}
pub fn add_generic_edge(&mut self, edge: TEdge) -> usize {
let idx = self.edges.len();
self.node_edges[edge.from_node()].push(idx);
self.reverse_node_edges[edge.to_node()].push(idx);
self.edges.push(edge);
idx
}
pub fn decompose(self) -> Vec<T> {
self.nodes
}
}
impl<T, TEdge: ReversibleEdge> GenericDirectedGraph<T, TEdge> {
pub fn reverse(mut self) -> Self {
(self.node_edges, self.reverse_node_edges) =
(self.reverse_node_edges, self.node_edges);
self.edges = self.edges
.into_iter()
.map(|edge| edge.reverse())
.collect();
self
}
}
impl<T, TEdge: Edge + From<(usize, usize)>> GenericDirectedGraph<T, TEdge> {
pub fn add_edge(&mut self, from: usize, to: usize) -> usize {
let edge = (from, to).into();
self.add_generic_edge(edge)
}
}
impl<T, TEdge: Edge> Graph for GenericDirectedGraph<T, TEdge> {
type Node = T;
type Edge = TEdge;
fn nodes(&self) -> impl Iterator<Item = &T> {
self.nodes.iter()
}
fn edges(&self) -> impl Iterator<Item = &TEdge> {
self.edges.iter()
}
fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.node_edges.get(node)
.map(|edges| edges.iter()
.map(|i| self.edges[*i].to_node()))
}
fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.reverse_node_edges.get(node)
.map(|edges| edges.iter()
.map(|i| self.edges[*i].from_node()))
}
fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
let reverse_edges = self.reverse_node_edges.get(node)?;
self.node_edges.get(node)
.map(|edges|
edges.iter()
.chain(reverse_edges.iter())
.map(|i| *i))
}
}
impl<T, TEdge: Edge> MutGraph for GenericDirectedGraph<T, TEdge> {
fn nodes_mut(&mut self) -> impl Iterator<Item = &mut Self::Node> {
self.nodes.iter_mut()
}
fn edges_mut(&mut self) -> impl Iterator<Item = &mut Self::Edge> {
self.edges.iter_mut()
}
}
/// An undirected graph that owns nodes and edges of
/// specific types given by generics.
pub struct GenericGraph<T, TEdge: Edge> {
nodes: Vec<T>,
edges: Vec<TEdge>,
node_neighbors: Vec<Vec<usize>>,
node_edges: Vec<Vec<usize>>,
adjacency_matrix: Option<Vec<Vec<Option<usize>>>>
}
impl<T, TEdge: Edge> GenericGraph<T, TEdge>
{
pub fn new(nodes: Vec<T>, adjacency_matrix: bool) -> Self {
let nodes_count = nodes.len();
Self {
nodes,
edges: vec![],
node_neighbors: vec![vec![]; nodes_count],
node_edges: vec![vec![]; nodes_count],
adjacency_matrix: if adjacency_matrix {
Some(vec![vec![None; nodes_count]; nodes_count])
} else {
None
}
}
}
pub fn add_generic_edge(&mut self, edge: TEdge) -> usize {
let idx = self.edges.len();
if let Some(adjacency_matrix) = self.adjacency_matrix.as_deref_mut() {
adjacency_matrix[edge.from_node()][edge.to_node()] = Some(idx);
adjacency_matrix[edge.to_node()][edge.from_node()] = Some(idx);
}
self.node_edges[edge.from_node()].push(idx);
self.node_edges[edge.to_node()].push(idx);
self.node_neighbors[edge.from_node()].push(edge.to_node());
self.node_neighbors[edge.to_node()].push(edge.from_node());
self.edges.push(edge);
idx
}
pub fn filter_edges(self, filter: impl Fn(usize, &TEdge) -> bool) -> Self {
let keep_edges = self.edges
.into_iter()
.enumerate()
.filter(|(i, e)| filter(*i, e))
.map(|(_, e)| e);
let mut new_graph = Self::new(self.nodes, self.adjacency_matrix.is_some());
for edge in keep_edges {
new_graph.add_generic_edge(edge);
}
new_graph
}
// NOTE: it's expected the edges will not reconnect, only the type will change.
// from_node() and to_node() should stay the same!
pub fn map_edges<TNewEdge: Edge>(self, map: impl Fn(TEdge) -> TNewEdge) -> GenericGraph<T, TNewEdge> {
GenericGraph::<T, TNewEdge> {
nodes: self.nodes,
edges: self.edges.into_iter().map(|edge| map(edge)).collect(),
node_neighbors: self.node_neighbors,
node_edges: self.node_edges,
adjacency_matrix: self.adjacency_matrix
}
}
pub fn map_nodes<TNewNode>(self, map: impl Fn(T) -> TNewNode) -> GenericGraph<TNewNode, TEdge> {
GenericGraph::<TNewNode, TEdge> {
nodes: self.nodes.into_iter().map(|node| map(node)).collect(),
edges: self.edges,
node_neighbors: self.node_neighbors,
node_edges: self.node_edges,
adjacency_matrix: self.adjacency_matrix
}
}
}
impl<T, TEdge: Edge> Graph for GenericGraph<T, TEdge> {
type Node = T;
type Edge = TEdge;
fn nodes(&self) -> impl Iterator<Item = &T> {
self.nodes.iter()
}
fn edges(&self) -> impl Iterator<Item = &TEdge> {
self.edges.iter()
}
fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.node_neighbors.get(node).map(|neighbors| neighbors.iter().map(|&node| node))
}
fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.node_neighbors.get(node).map(|neighbors| neighbors.iter().map(|&node| node))
}
fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.node_edges.get(node).map(|edges| edges.iter().map(|&edge| edge))
}
fn has_edge(&self, from: usize, to: usize) -> bool {
if let Some(adjacency_matrix) = &self.adjacency_matrix {
adjacency_matrix[from][to].is_some()
} else {
self.edges()
.any(|edge| edge.from_node() == from && edge.to_node() == to)
}
}
fn get_edge_between(&self, from: usize, to: usize) -> Option<usize> {
if let Some(adjacency_matrix) = &self.adjacency_matrix {
adjacency_matrix[from][to]
} else {
self.edges_of_idxs(from)
.map(|edges| edges
.filter(|&edge| self.edge(edge).unwrap().to_node() == to)
.next())
.flatten()
}
}
}
impl<T, TEdge: Edge> MutGraph for GenericGraph<T, TEdge> {
fn nodes_mut(&mut self) -> impl Iterator<Item = &mut Self::Node> {
self.nodes.iter_mut()
}
fn edges_mut(&mut self) -> impl Iterator<Item = &mut Self::Edge> {
self.edges.iter_mut()
}
}
pub struct ReversedEdge<'a, T: Edge> {
edge: &'a T
}
impl<'a, T: Edge> ReversedEdge<'a, T> {
pub fn new(edge: &'a T) -> Self {
Self {
edge
}
}
}
impl<'a, T: Edge> Edge for ReversedEdge<'a, T> {
fn from_node(&self) -> usize {
self.edge.to_node()
}
fn to_node(&self) -> usize {
self.edge.from_node()
}
}
/// A view on a graph that reverses all its
/// edges.
pub struct ReversedGraph<'a, T: Graph> {
graph: &'a T,
edges: Vec<ReversedEdge<'a, T::Edge>>
}
impl<'a, T: Graph> ReversedGraph<'a, T> {
pub fn new(graph: &'a T) -> Self {
Self {
graph,
edges: graph.edges()
.map(|edge| ReversedEdge::new(edge))
.collect()
}
}
}
impl<'a, T: Graph> Graph for ReversedGraph<'a, T> {
type Node = T::Node;
type Edge = ReversedEdge<'a, T::Edge>;
fn nodes(&self) -> impl Iterator<Item = &Self::Node> {
self.graph.nodes()
}
fn edges(&self) -> impl Iterator<Item = &Self::Edge> {
self.edges.iter()
}
fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.graph.reverse_neighbor_idxs(node)
}
fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.graph.neighbor_idxs(node)
}
fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
self.graph.edges_of_idxs(node)
}
}
/// Make a search out of the starting node, visiting every
/// node that can be visited from it. Denoting the
/// distances between the nodes and the parents that visited the
/// given node for reconstructing the shortest path.
pub fn breadth_first_search<TGraph>(
graph: &TGraph,
starting_node: usize,
) -> (Vec<Option<Distance>>, Vec<usize>)
where
TGraph: Graph
{
let nodes_len = graph.nodes().count();
let mut distances = vec![None; nodes_len];
let mut visited = vec![false; nodes_len];
let mut parents = vec![0; nodes_len];
let mut node_queue = VecDeque::with_capacity(nodes_len / 4);
node_queue.push_back(starting_node);
distances[starting_node] = Some(0);
visited[starting_node] = true;
while let Some(current) = node_queue.pop_front() {
for neighbor in graph.neighbor_idxs(current).unwrap() {
if visited[neighbor] {
continue
}
visited[neighbor] = true;
distances[neighbor] = Some(distances[current].unwrap() + 1);
parents[neighbor] = current;
node_queue.push_back(neighbor);
}
}
(distances, parents)
}
/// Like breadth first search, but only look if a node
/// is reachable, do not figure out the distance, do not
/// figure out the shortest path.
pub fn breadth_first_reachable<TGraph>(
graph: &TGraph,
starting_node: usize,
) -> Vec<bool>
where
TGraph: Graph
{
let nodes_len = graph.nodes().count();
let mut visited = vec![false; nodes_len];
let mut node_queue = VecDeque::with_capacity(nodes_len / 4);
node_queue.push_back(starting_node);
visited[starting_node] = true;
while let Some(current) = node_queue.pop_front() {
for neighbor in graph.neighbor_idxs(current).unwrap() {
if visited[neighbor] {
continue
}
visited[neighbor] = true;
node_queue.push_back(neighbor);
}
}
visited
}
/// Figure out distance from each node to each node.
/// In case the node is unreachable, Distance::MAX is
/// in the matrix.
pub fn floyd_warshall<TNode, TEdge, TGraph>(
graph: &TGraph
) -> Vec<Vec<Distance>>
where
TEdge: Edge,
TGraph: Graph<Node = TNode, Edge = TEdge>
{
let nodes = graph.nodes().count();
let mut distances = vec![vec![Distance::MAX; nodes]; nodes];
for edge in graph.edges() {
distances[edge.from_node()][edge.to_node()] = 1;
}
for v in 0..nodes {
distances[v][v] = 0;
}
for k in 0..nodes {
for i in 0..nodes {
for j in 00..nodes {
if distances[i][k] == Distance::MAX || distances[k][j] == Distance::MAX {
continue;
}
if distances[i][j] > distances[i][k] + distances[k][j] {
distances[i][j] = distances[i][k] + distances[k][j];
}
}
}
}
distances
}
/// A generic representation of a minimum spanning
/// tree for cases where the graph might not be
/// fully connected, and thus it is possible the minimum
/// spanning tree will not be a single component.
#[derive(Clone, Debug, PartialEq)]
pub struct MinimumSpanningTree {
// What component does node i belong to?
pub components: UnionFind,
// Edge indices in the tree
pub edges: Vec<usize>,
}
impl MinimumSpanningTree {
pub fn nodes_count(&self) -> usize {
self.components.len()
}
pub fn components_count(&self) -> usize {
self.nodes_count() - self.edges.len()
}
}
/// Use the kruskal algorithm for finding the minimum spanning tree.
/// Take only edges filtered by selector as candidates for the spanning tree.
/// Initial minimum spanning tree can be passed, that should usually be a result
/// of prior run of this function with a different selector.
pub fn minimal_spanning_tree_kruskal<'a, TNode, TWeight, TEdge, TGraph>(
graph: &TGraph,
initial: Option<MinimumSpanningTree>,
selector: impl Fn(&TEdge) -> bool
) -> MinimumSpanningTree
where
TWeight: PartialOrd,
TEdge: Edge + WeightedEdge<Cost = TWeight>,
TGraph: Graph<Node = TNode, Edge = TEdge>
{
// let separate_new_edges = initial.is_some();
let nodes = graph.nodes().count();
let mut initial_edge_selected = vec![false; graph.edges().count()];
let mut current = if let Some(initial) = initial {
for edge in &initial.edges {
initial_edge_selected[*edge] = true;
}
initial
} else {
MinimumSpanningTree {
components: UnionFind::make_set(nodes),
edges: Vec::with_capacity(nodes - 1),
}
};
let mut remaining_edges = graph.edges()
.enumerate()
.filter(|(i, e)| !initial_edge_selected[*i] && selector(e))
.collect::<Vec<_>>();
remaining_edges.sort_by(
|a, b| a.1.cost().partial_cmp(&b.1.cost()).unwrap_or(Ordering::Less)
);
for (i, edge) in remaining_edges {
// 1. does the edge connect two components?
let (root_a, root_b) = {
if current.components_count() == 1 {
break;
}
let root_a = current.components.find(edge.from_node());
let root_b = current.components.find(edge.to_node());
if root_a == root_b {
continue;
}
(root_a, root_b)
};
// 2. if so, use it and connect them
current.components.union(root_a, root_b);
current.edges.push(i);
}
current
}