gtrs/src/graph.rs

use std::{
    collections::{BTreeMap, HashMap, HashSet, VecDeque},
    fmt::Debug,
    hash::Hash,
};

pub trait VertexType: Eq + Hash + Ord + Clone + Debug {}

impl<T> VertexType for T where T: Eq + Hash + Ord + Clone + Debug {}

pub struct Graph<V> {
    pub vertices: HashSet<V>,
    pub edges: HashMap<V, HashSet<V>>,
}

impl<V> Graph<V>
where
    V: VertexType,
{
    /// get the number of vertices of the graph
    pub fn vertex_count(&self) -> usize {
        self.vertices.len()
    }

    /// get the number of edges of the graph
    pub fn edge_count(&self) -> usize {
        self.edges.values().map(HashSet::len).sum()
    }

    /// check if an edge is contained in the graph
    /// (only checks the provided direction)
    pub fn has_edge(&self, edge: (&V, &V)) -> bool {
        let (from, to) = edge;
        self.edges
            .get(from)
            .map_or(false, |v| v.iter().any(|x| x == to))
    }

    /// check if an edge (a, b) and its' reverse (b, a) are both contained in the graph
    pub fn has_bidirectional_edge(&self, edge: (&V, &V)) -> bool {
        let (a, b) = edge;
        self.has_edge((a, b)) && self.has_edge((b, a))
    }

    /// check if a vertex is contained in the graph
    pub fn has_vertex(&self, vertex: &V) -> bool {
        self.vertices.contains(vertex)
    }

    /// get a slice containing all neighbors of the edge
    pub fn neighbors(&self, vertex: &V) -> Vec<&V> {
        self.edges
            .get(vertex)
            .map(|es| es.iter().collect())
            .unwrap_or_default()
    }

    /// find a path between two nodes using breadth-first-search.
    /// this finds a path with the least possible edges.
    pub fn find_path_bfs(&self, from: &V, to: &V) -> Option<Vec<V>> {
        let mut q = VecDeque::with_capacity(self.vertices.len());
        let mut visited = HashSet::with_capacity(self.vertices.len());

        q.push_back(vec![from]);
        visited.insert(from);

        while let Some(mut path) = q.pop_front() {
            let current = path.last().unwrap();

            for neighbor in self.neighbors(current) {
                if neighbor == to {
                    return path
                        .into_iter()
                        .cloned()
                        .chain([to.clone()])
                        .collect::<Vec<_>>()
                        .into();
                }

                path.push(neighbor);
                q.push_back(path.clone());
                path.pop();
            }
        }

        None
    }

    /// find a path between two nodes using depth-first-search.
    ///
    /// this is short-circuiting and therefore does not guarantee
    /// the found path to contain the least possible edges.
    pub fn find_path_dfs(&self, from: &V, to: &V) -> Option<Vec<V>> {
        let mut q = Vec::with_capacity(self.vertices.len());

        q.push(vec![from]);

        while let Some(mut path) = q.pop() {
            let &current = path.last().unwrap();

            if current == to {
                return path.into_iter().cloned().collect::<Vec<_>>().into();
            }

            for neighbor in self.neighbors(current).iter().rev() {
                if path.contains(neighbor) {
                    continue;
                }

                path.push(neighbor);
                q.push(path.clone());
                path.pop();
            }
        }

        None
    }

    pub fn iter(&self) -> impl Iterator<Item = &V> {
        self.vertices.iter()
    }
}

impl<V> IntoIterator for Graph<V>
where
    V: VertexType,
{
    type Item = V;
    type IntoIter = IntoIter<V>;

    fn into_iter(self) -> Self::IntoIter {
        let items = self.vertices.into_iter().collect::<Vec<_>>().clone();
        IntoIter { items, i: 0 }
    }
}

pub struct IntoIter<V> {
    items: Vec<V>,
    i: usize,
}

impl<V> Iterator for IntoIter<V>
where
    V: VertexType,
{
    type Item = V;
    fn next(&mut self) -> Option<Self::Item> {
        match self.items.get(self.i) {
            None => None,
            item => {
                self.i += 1;
                item.cloned()
            }
        }
    }
}

pub struct WeightedGraph<V> {
    pub vertices: HashSet<V>,
    pub edges: HashMap<V, HashSet<(u64, V)>>,
}

impl<V> WeightedGraph<V>
where
    V: VertexType,
{
    /// create a weighted graph from a set of vertices and weighted edges.
    pub fn new<I, J>(vertices: I, edges: J) -> Self
    where
        I: IntoIterator<Item = V>,
        J: IntoIterator<Item = (V, V, u64)>,
    {
        let vertices = vertices.into_iter().collect();
        let mut parsed_edges = HashMap::new();
        for (from, to, c) in edges {
            parsed_edges
                .entry(from)
                .and_modify(|e: &mut HashSet<(u64, V)>| {
                    e.insert((c, to.clone()));
                })
                .or_insert_with(|| HashSet::from([(c, to)]));
        }

        Self {
            vertices,
            edges: parsed_edges,
        }
    }

    /// get the number of vertices of the graph
    pub fn vertex_count(&self) -> usize {
        self.vertices.len()
    }

    /// get the number of edges of the graph
    pub fn edge_count(&self) -> usize {
        self.edges.values().map(HashSet::len).sum()
    }

    /// check if an edge is contained in the graph
    /// (only checks the provided direction)
    pub fn has_edge(&self, edge: (&V, &V)) -> bool {
        let (from, to) = edge;
        self.edges
            .get(from)
            .map_or(false, |v| v.iter().any(|(_, x)| x == to))
    }

    /// check if an edge (a, b) and its' reverse (b, a) are both contained in the graph
    pub fn has_bidirectional_edge(&self, edge: (&V, &V)) -> bool {
        let (a, b) = edge;
        self.has_edge((a, b)) && self.has_edge((b, a))
    }

    /// check if a vertex is contained in the graph
    pub fn has_vertex(&self, vertex: &V) -> bool {
        self.vertices.contains(vertex)
    }

    /// get a slice containing all neighbors of the edge
    pub fn neighbors(&self, vertex: &V) -> Vec<&(u64, V)> {
        self.edges
            .get(vertex)
            .map(|es| es.iter().collect())
            .unwrap_or_default()
    }

    /// find a path between two nodes using breadth-first-search.
    /// this finds a path with the least possible edges.
    ///
    /// does not take into account edge weights.
    pub fn find_path_bfs(&self, from: &V, to: &V) -> Option<Vec<V>> {
        let mut q = VecDeque::with_capacity(self.vertices.len());
        let mut visited = HashSet::with_capacity(self.vertices.len());

        q.push_back(vec![from]);
        visited.insert(from);

        while let Some(mut path) = q.pop_front() {
            let current = path.last().unwrap();

            for (_, neighbor) in self.neighbors(current) {
                if neighbor == to {
                    return path
                        .into_iter()
                        .cloned()
                        .chain([to.clone()])
                        .collect::<Vec<_>>()
                        .into();
                }

                path.push(neighbor);
                q.push_back(path.clone());
                path.pop();
            }
        }

        None
    }

    /// find a path between two nodes using depth-first-search.
    ///
    /// this is short-circuiting and therefore does not guarantee
    /// the found path to contain the least possible edges.
    ///
    /// does not take into account edge weights.
    pub fn find_path_dfs(&self, from: &V, to: &V) -> Option<Vec<V>> {
        let mut q = Vec::with_capacity(self.vertices.len());

        q.push(vec![from]);

        while let Some(mut path) = q.pop() {
            let &current = path.last().unwrap();

            if current == to {
                return path.into_iter().cloned().collect::<Vec<_>>().into();
            }

            for (_, neighbor) in self.neighbors(current).iter().rev() {
                if path.contains(&neighbor) {
                    continue;
                }

                path.push(neighbor);
                q.push(path.clone());
                path.pop();
            }
        }

        None
    }

    pub fn find_path_dijkstra(&self, from: &V, to: &V) -> Option<Vec<V>> {
        use std::cmp::Reverse as Rev;
        use std::collections::BinaryHeap;

        let mut dist: BTreeMap<_, _> = self
            .vertices
            .iter()
            .zip(std::iter::repeat(u64::MAX))
            .collect();

        let mut prev: HashMap<_, _> = self.vertices.iter().zip(std::iter::repeat(None)).collect();
        let mut q: BinaryHeap<_> = self.vertices.iter().map(Rev).collect();
        dist.entry(from).and_modify(|c| *c = 0);

        let mut i = 0;
        while !q.is_empty() {
            let u = q.iter().min_by_key(|v| dist[v.0]).unwrap().0.clone();

            if &u == to {
                let mut s = VecDeque::new();
                let mut u = Some(to.clone());
                if prev[&u.clone()?].is_some() || u.clone()? == *from {
                    while u.is_some() {
                        s.push_front(u.clone()?);
                        u = prev[&u?].clone();
                    }
                    return Some(s.iter().cloned().collect());
                }

                return None;
            }

            q.retain(|v| *v != Rev(&u)); // remove u
            i += 1;
            if i > 5 {
                panic!();
            }

            for (weight, v) in self.neighbors(&u) {
                let alt = dist[&u] + weight;
                if alt < dist[&v] {
                    *dist.get_mut(v)? = alt;
                    *prev.get_mut(v)? = Some(u.clone());
                }
            }
        }

        None
    }

    pub fn iter(&self) -> impl Iterator<Item = &V> {
        self.vertices.iter()
    }
}

#[macro_export]
macro_rules! graph {
    ( $( $v:tt : $( $e:tt ),* );* $(;)? ) => {{
        let mut edges = std::collections::HashMap::new();
        let mut vertices = std::collections::HashSet::new();

        $(
            vertices.insert($v); // insert vertex
            let v_neighbors = std::collections::HashSet::from([
                $( $e ),*
            ]);

            edges.entry($v).or_insert(v_neighbors); // initialize its edges

            // insert vertices with no outgoing edges
            $( vertices.insert($e); )*
        )*

        Graph { vertices, edges }
    }};
}