Graph.hpp

 
#pragma once
 
#include <functional>
#include <iostream>
#include <numeric>
#include <queue>
#include <stack>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
 
namespace dlinear {
 
enum class VisitResult {
  CONTINUE,  
  SKIP,      
  STOP,      
};
 
template <class T, class W>
struct EdgeHash_ {
  size_t operator()(const std::pair<T, W>& lhs) const { return std::hash<T>{}(lhs.first); }
};
 
template <class T, class W>
struct EdgeEqual_ {
  bool operator()(const std::pair<T, W>& lhs, const std::pair<T, W>& rhs) const {
    return std::equal_to<T>{}(lhs.first, rhs.first);
  }
};
 
// using T = int;
// using W = double;
// using EdgeHash = EdgeHash_<T, W>;
// using EdgeEqual = EdgeEqual_<T, W>;
 
template <class T, std::default_initializable W, class EdgeHash = EdgeHash_<T, W>, class EdgeEqual = EdgeEqual_<T, W>>
class Graph {
 public:
  using Edge = std::pair<T, W>;
  using AdjSet = std::unordered_set<Edge, EdgeHash, EdgeEqual>;
  using VisitFunc = std::function<VisitResult(const T&, const T&, const W&)>;  
  using PathsFunc = std::function<VisitResult(
      std::vector<T>&)>;  
 
  void AddEdge(const T& u, const T& v, bool bidirectional = true) {
    adj_list_[u].emplace(v, 1);
    adj_list_[v];  // Ensure the ending vertex exists
    if (bidirectional) adj_list_[v].emplace(u, 1);
  }
 
  bool AddEdge(const T& u, const T& v, const W& weight, bool bidirectional = true) {
    bool updated = false;
    const auto [it, inserted] = adj_list_[u].emplace(v, weight);
    adj_list_[v];  // Ensure the ending vertex exists
    if (!inserted && it->second != weight) {
      adj_list_.at(u).erase(it);
      adj_list_.at(u).emplace(v, weight);
      updated = true;
    }
    if (bidirectional) {
      const auto [b_it, b_inserted] = adj_list_[v].emplace(u, 1 / weight);
      if (!b_inserted && b_it->second != 1 / weight) {
        adj_list_.at(v).erase(b_it);
        adj_list_.at(v).emplace(u, 1 / weight);
        updated = true;
      }
    }
    return updated;
  }
 
  void AddVertex(const T& v) {
    if (adj_list_.find(v) == adj_list_.end()) adj_list_.try_emplace(v);
  }
 
  bool HasEdge(const T& u, const T& v) const {
    return adj_list_.find(u) != adj_list_.cend() && adj_list_.at(u).find({v, W{}}) != adj_list_.at(u).cend();
  }
  bool HasEdges(const T& u) const { return adj_list_.find(u) != adj_list_.cend() && !adj_list_.at(u).empty(); }
 
  const W* GetEdgeWeight(const T& u, const T& v) const {
    if (auto it = adj_list_.find(u); it != adj_list_.cend()) {
      if (auto it2 = adj_list_.at(u).find({v, W{}}); it2 != adj_list_.at(u).cend()) {
        return &it2->second;
      }
    }
    return nullptr;
  }
 
  const std::unordered_map<T, AdjSet>& adj_list() const { return adj_list_; }
 
  void RemoveEdge(const T& u, const T& v, bool bidirectional = true) {
    adj_list_[u].erase({v, W{}});
    if (bidirectional) adj_list_[v].erase({u, W{}});
  }
 
  void RemoveVertex(const T& v) {
    adj_list_.erase(v);
    for (auto& [node, edges] : adj_list_) edges.erase({v, W{}});
  }
 
  void Clear() { adj_list_.clear(); }
  void ClearEdges() {
    for (auto& [node, edges] : adj_list_) edges.clear();
  }
 
  [[nodiscard]] size_t Size() const {
    return std::accumulate(adj_list_.begin(), adj_list_.end(), 0,
                           [](size_t acc, const auto& pair) { return acc + pair.second.size(); });
  }
 
  [[nodiscard]] bool IsEmpty() const { return adj_list_.empty(); }
 
  void DFS(const T& start, const VisitFunc& func) {
    std::unordered_set<T> visited{};
    DFS(start, func, visited);
  }
  void DFS(const T& start, const VisitFunc& func, std::unordered_set<T>& visited) {
    // If the starting vertex is not in the graph, return
    if (adj_list_.find(start) == adj_list_.end()) return;
 
    visited.reserve(adj_list_.size());
    std::stack<T> stack;
    std::unordered_map<T, std::pair<T, const W*>> edges;
 
    edges.emplace(start, std::make_pair(start, &zero_));
    stack.push(start);
    while (!stack.empty()) {
      const T current = std::move(stack.top());
      stack.pop();
      if (visited.find(current) != visited.end()) continue;
      visited.insert(current);
      const VisitResult res = func(edges.at(current).first, current, *edges.at(current).second);
      if (res == VisitResult::STOP) return;
      if (res == VisitResult::SKIP || adj_list_.find(current) == adj_list_.end()) continue;
      for (auto adj_it = adj_list_.at(current).begin(); adj_it != adj_list_.at(current).end(); ++adj_it) {
        const auto& [adj_vertex, weight] = *adj_it;
        if (visited.find(adj_vertex) != visited.end()) continue;
        stack.push(adj_vertex);
        edges.insert_or_assign(adj_vertex, std::make_pair(current, &weight));
      }
    }
  }
 
  void BFS(const T& start, const VisitFunc& func) {
    std::unordered_set<T> visited{};
    BFS(start, func, visited);
  }
  void BFS(const T& start, const VisitFunc& func, std::unordered_set<T>& visited) {
    // If the starting vertex is not in the graph, return
    if (adj_list_.find(start) == adj_list_.end()) return;
 
    visited.reserve(adj_list_.size());
    std::queue<T> queue;
    std::unordered_map<T, std::pair<T, const W*>> edges;
 
    visited.insert(start);
    edges.emplace(start, std::make_pair(start, &zero_));
    queue.push(start);
    while (!queue.empty()) {
      const VisitResult res = func(edges.at(queue.front()).first, queue.front(), *edges.at(queue.front()).second);
      if (res == VisitResult::STOP) return;
      if (res == VisitResult::SKIP || adj_list_.find(queue.front()) == adj_list_.end()) {
        queue.pop();
        continue;
      }
      for (auto adj_it = adj_list_.at(queue.front()).begin(); adj_it != adj_list_.at(queue.front()).end(); ++adj_it) {
        const auto& [adj_vertex, weight] = *adj_it;
        if (visited.find(adj_vertex) != visited.end()) continue;
        visited.insert(adj_vertex);
        queue.push(adj_vertex);
        edges.insert_or_assign(adj_vertex, std::make_pair(queue.front(), &weight));
      }
      queue.pop();
    }
  }
 
  void AllPaths(const T& start, const T& end, const PathsFunc& func) {
    std::unordered_set<T> visited;
    AllPaths(start, end, func, visited);
  }
  void AllPaths(const T& start, const T& end, const PathsFunc& func, std::unordered_set<T>& visited) {
    if (adj_list_.find(start) == adj_list_.end() || adj_list_.find(end) == adj_list_.end() ||
        adj_list_.find(start) == adj_list_.find(end))
      return;
    std::stack<T> stack;
    std::unordered_map<T, typename std::unordered_set<Edge>::iterator> iterators;
    std::vector<T> path;
 
    iterators.reserve(adj_list_.size());
    visited.reserve(adj_list_.size());
    path.reserve(adj_list_.size());
 
    stack.push(start);
    visited.insert(start);
    path.push_back(start);
    iterators[start] = adj_list_[start].begin();
 
    while (!stack.empty()) {
      const T current = stack.top();
 
      if (current == end) {
        if (func(path) != VisitResult::CONTINUE) return;
        stack.pop();
        visited.erase(current);
        path.pop_back();
        continue;
      }
 
      auto& it = iterators.at(current);
      if (it == adj_list_.at(current).end()) {
        stack.pop();
        visited.erase(current);
        path.pop_back();
        continue;
      }
 
      const T next = it->first;
      ++it;
      if (visited.find(next) == visited.end()) {
        stack.push(next);
        visited.insert(next);
        path.push_back(next);
        iterators.insert_or_assign(next, adj_list_.at(next).begin());
      }
    }
  }
 
 private:
  const W zero_{};
  std::unordered_map<T, AdjSet> adj_list_;  
};
 
template <class T, class E>
std::ostream& operator<<(std::ostream& os, const Graph<T, E>& s) {
  os << "Graph{";
  for (const auto& [vertex, edges] : s.adj_list()) {
    os << "| " << vertex << " | -> ";
    for (const auto& [adj_vertex, weight] : edges) os << adj_vertex << "(" << weight << "), ";
    os << " - ";
  }
  return os << "}";
}
 
}  // namespace dlinear