2023 graph: return edges between components from tarjan_scc

2023/src/day20.hpp

@@ -18,10 +18,10 @@
 #include <memory>        // for unique_ptr, make_unique
 #include <numeric>       // for lcm
 #include <queue>         // for queue
+#include <set>           // for set (tarjan_scc)
 #include <sstream>       // for istringstream
 #include <string>        // for string, getline
 #include <unordered_map> // for unordered_map
-#include <unordered_set> // for unordered_set
 #include <utility>       // for move, forward
 #include <vector>        // for vector
 // IWYU pragma: no_include <functional>  // for hash (unordered_map)
@@ -116,7 +116,7 @@ struct CycleInfo {
 class MessageBus {
     std::unordered_map<std::string, std::unique_ptr<ModuleBase>> modules;
     std::queue<Message> msg_queue;
-    std::vector<std::unordered_set<std::string>> components{};
+    std::vector<std::vector<std::string>> components{};
     MessageCounter counter;
 
     bool _rx_activated = false;
@@ -338,7 +338,7 @@ void MessageBus::identify_components() {
     const auto get_neighbors = [this](const std::string &name) {
         return modules.at(name)->outputs;
     };
-    components = aoc::graph::tarjan_scc(root, get_neighbors);
+    components = aoc::graph::tarjan_scc(root, get_neighbors).first;
 
     // store component ids in each module
     for (std::size_t i = 0; i < components.size(); ++i) {

2023/src/graph_traversal.hpp

@@ -10,8 +10,10 @@
 #ifndef GRAPH_TRAVERSAL_HPP_56T9ZURK
 #define GRAPH_TRAVERSAL_HPP_56T9ZURK
 
+#include "lib.hpp"           // for DEBUG
 #include "util/concepts.hpp" // for Hashable, any_iterable_collection, same_as_any
 #include <algorithm>         // for min, reverse
+#include <cassert>           // for assert
 #include <concepts>          // for same_as, integral
 #include <functional>        // for function, greater
 #include <map>               // for map
@@ -282,66 +284,102 @@ namespace detail {
 struct tarjan_entry {
     int index = -1;
     int lowlink = -1;
-    bool onStack = false;
+    int component_id = -1;
+    std::vector<int> pending_edges;
 };
 } // namespace detail
 
 /**
  * Find strongly connected components of a directed graph.
  *
- * Components are returned in topological order.
+ * Components are returned in topological order, along with a set of the
+ * directed edges between the components.
  */
 template <class Key, detail::GetNeighbors<Key> GetNeighbors>
-std::vector<detail::maybe_unordered_set<Key>>
+std::pair<std::vector<std::vector<Key>>, std::set<std::pair<int, int>>>
 tarjan_scc(const Key &source, GetNeighbors &&get_neighbors) {
     int index = 0;
     std::stack<Key> S{};
     std::vector<detail::maybe_unordered_set<Key>> components{};
 
     detail::maybe_unordered_map<Key, detail::tarjan_entry> entries;
+    std::set<std::pair<int, int>> component_links;
 
-    const auto strongconnect = [&](const Key &v, auto &rec) -> void {
+    const auto strongconnect =
+        [&get_neighbors, &index, &S, &components, &entries,
+         &component_links](const Key &v, auto &rec) -> detail::tarjan_entry & {
         detail::tarjan_entry &v_entry = entries[v];
         v_entry.index = index;
         v_entry.lowlink = index;
         ++index;
         S.push(v);
-        v_entry.onStack = true;
+        assert(v_entry.component_id == -1);
 
         for (const Key &w : get_neighbors(v)) {
-            if (!entries.contains(w)) {
+            detail::tarjan_entry *w_entry = nullptr;
+            auto it = entries.find(w);
+            if (it == entries.end()) {
                 // Successor w has not yet been visited; recurse on it
-                rec(w, rec);
-                detail::tarjan_entry &w_entry = entries.at(w);
-                v_entry.lowlink = std::min(v_entry.lowlink, w_entry.lowlink);
-                continue;
-            }
-            detail::tarjan_entry &w_entry = entries.at(w);
-            if (w_entry.onStack) {
-                // Successor w is in stack S and hence in the current SCC
-                // If w is not on stack, then (v, w) is an edge pointing to an
-                // SCC already found and must be ignored
+                w_entry = &rec(w, rec);
+                v_entry.lowlink = std::min(v_entry.lowlink, w_entry->lowlink);
+            } else {
+                w_entry = &it->second;
+                if (it->second.component_id == -1) {
+                    // Successor w is in stack S and hence in the current SCC
+                    // If w is not on stack, then (v, w) is an edge pointing to
+                    // an SCC already found and must be ignored
 
-                v_entry.lowlink = std::min(v_entry.lowlink, w_entry.index);
+                    v_entry.lowlink = std::min(v_entry.lowlink, w_entry->index);
+                }
+            }
+            if (w_entry->component_id != -1) {
+                v_entry.pending_edges.push_back(w_entry->component_id);
             }
         }
         // If v is a root node, pop the stack and generate an SCC
         if (v_entry.lowlink == v_entry.index) {
+            int component_id = components.size();
             components.emplace_back();
-            Key w;
+            Key w = S.top();
             do {
                 w = S.top();
                 S.pop();
-                entries.at(w).onStack = false;
-                components.back().insert(w);
+                auto &entry = entries.at(w);
+                entry.component_id = component_id;
+                for (int x : entry.pending_edges) {
+                    component_links.emplace(component_id, x);
+                }
+                components.back().push_back(w);
             } while (w != v);
         }
+        return v_entry;
     };
 
     strongconnect(source, strongconnect);
 
+    // check edges
+    if constexpr (aoc::DEBUG) {
+        // reconstruct the links between components manually
+        std::set<std::pair<int, int>> reconstructed_links;
+        for (const auto &[v, v_entry] : entries) {
+            for (const Key &w : get_neighbors(v)) {
+                const auto &w_entry = entries.at(w);
+                if (v_entry.component_id != w_entry.component_id) {
+                    reconstructed_links.emplace(v_entry.component_id,
+                                                w_entry.component_id);
+                }
+            }
+        }
+        assert(component_links == reconstructed_links);
+    }
+
     std::ranges::reverse(components);
-    return components;
+    std::set<std::pair<int, int>> reversed_links;
+    for (const auto &[v_id, w_id] : component_links) {
+        reversed_links.emplace(components.size() - 1 - v_id,
+                               components.size() - 1 - w_id);
+    }
+    return {std::move(components), std::move(reversed_links)};
 }
 
 /**

2023/src/test00_graph.cpp

@@ -9,23 +9,29 @@
 #include "unit_test/unit_test.hpp"
 
 #include "graph_traversal.hpp"
+#include <algorithm>     // for sort
 #include <cstddef>       // for size_t
+#include <set>           // for set
 #include <unordered_map> // for unordered_map
-#include <unordered_set> // for unordered_set
+#include <utility>       // for pair
 #include <vector>        // for vector
-// IWYU pragma: no_include <algorithm>  // for copy
 
 namespace aoc::graph::test {
 
 std::size_t test_tarjan_scc() {
     unit_test::PureTest test(
         "aoc::graph::tarjan_scc",
         +[](const std::unordered_map<int, std::vector<int>> &adj,
-            int root) -> std::vector<std::unordered_set<int>> {
+            int root) -> std::pair<std::vector<std::vector<int>>,
+                                   std::set<std::pair<int, int>>> {
             const auto get_neighbors = [&adj](int key) -> std::vector<int> {
                 return adj.at(key);
             };
-            return aoc::graph::tarjan_scc(root, get_neighbors);
+            auto result = aoc::graph::tarjan_scc(root, get_neighbors);
+            for (auto &component : result.first) {
+                std::sort(component.begin(), component.end());
+            }
+            return result;
         });
 
     // example from Tarjan's original 1972 paper, "Depth-first search and
@@ -38,7 +44,7 @@ std::size_t test_tarjan_scc() {
           {6, {}},
           {7, {4, 6}},
           {8, {1, 7}}},
-         1, {{1, 2, 8}, {3, 4, 5, 7}, {6}});
+         1, {{{1, 2, 8}, {3, 4, 5, 7}, {6}}, {{0, 1}, {1, 2}}});
     // example from Wikipedia:
     // https://commons.wikimedia.org/wiki/File:Tarjan%27s_Algorithm_Animation.gif
     // numbering:
@@ -52,7 +58,9 @@ std::size_t test_tarjan_scc() {
           {6, {8}},
           {7, {6}},
           {8, {7}}},
-         1, {{1}, {3, 5}, {2, 4}, {6, 7, 8}});
+         1,
+         {{{1}, {3, 5}, {2, 4}, {6, 7, 8}},
+          {{0, 1}, {0, 2}, {1, 2}, {1, 3}, {2, 3}}});
 
     return test.done(), test.num_failed();
 }