Separate spfa and deep_spfa, add spfa tree diameter test

Author: adamant-pwn

cp-algo/graph/shortest_path.hpp

@@ -16,7 +16,7 @@ namespace cp_algo::graph {
             : dist(n, inf), pre(n) {}
     };
 
-    struct dijkstra_context : shortest_path_context {
+    struct dijkstra_context: shortest_path_context {
         struct que_t {
             int64_t dist;
             node_index v;
@@ -44,23 +44,44 @@ namespace cp_algo::graph {
         }
     };
 
-    struct spfa_context : shortest_path_context {
+    struct spfa_context: shortest_path_context {
         std::queue<node_index> que;
         std::vector<char> flags;
         static constexpr char in_queue = 1;
         static constexpr char invalidated = 2;
+
+        spfa_context(int n) : shortest_path_context(n), flags(n) {}
+
+        void push(node_index, node_index v) {
+            if (!(flags[v] & in_queue)) {
+                que.push(v);
+                flags[v] |= in_queue;
+            }
+        }
+
+        std::optional<node_index> next_node() {
+            while (!que.empty()) {
+                node_index v = que.front();
+                que.pop();
+                flags[v] &= ~in_queue;
+                if (!(flags[v] & invalidated)) {
+                    return v;
+                }
+            }
+            return std::nullopt;
+        }
+    };
+
+    struct deep_spfa_context: spfa_context {
         std::vector<std::basic_string<node_index>> dependents;
 
-        spfa_context(int n) : shortest_path_context(n), flags(n), dependents(n) {}
+        deep_spfa_context(int n) : spfa_context(n), dependents(n) {}
 
         void push(node_index u, node_index v) {
             invalidate_subtree(v);
             dependents[u].push_back(v);
             flags[v] &= ~invalidated;
-            if (!(flags[v] & in_queue)) {
-                que.push(v);
-                flags[v] |= in_queue;
-            }
+            spfa_context::push(u, v);
         }
 
         void invalidate_subtree(node_index v) {
@@ -69,6 +90,7 @@ namespace cp_algo::graph {
                 node_index u = to_invalidate.back();
                 to_invalidate.pop_back();
                 flags[u] |= invalidated;
+                flags[u] &= ~in_queue;
                 for (auto dep: dependents[u]) {
                     if (pre[dep].u == u) {
                         to_invalidate.push_back(dep);
@@ -77,18 +99,6 @@ namespace cp_algo::graph {
                 dependents[u].clear();
             }
         }
-
-        std::optional<node_index> next_node() {
-            while (!que.empty()) {
-                node_index v = que.front();
-                que.pop();
-                flags[v] &= ~in_queue;
-                if (!(flags[v] & invalidated)) {
-                    return v;
-                }
-            }
-            return std::nullopt;
-        }
     };
 
     template<typename Context, weighted_graph_type graph>
@@ -119,14 +129,18 @@ namespace cp_algo::graph {
     shortest_path_context spfa(graph const& g, node_index s) {
         return sssp_impl<spfa_context>(g, s);
     }
+    template<weighted_graph_type graph>
+    shortest_path_context deep_spfa(graph const& g, node_index s) {
+        return sssp_impl<deep_spfa_context>(g, s);
+    }
 
     template<weighted_graph_type graph>
     shortest_path_context single_source_shortest_path(graph const& g, node_index s) {
         bool negative_edges = false;
         for (auto e: g.edges()) {
             negative_edges |= e.w < 0;
         }
-        return negative_edges ? spfa(g, s) : dijkstra(g, s);
+        return negative_edges ? deep_spfa(g, s) : dijkstra(g, s);
     }
 
     std::vector<edge_index> recover_path(auto const& pre, node_index s, node_index t) {

verify/graph/tree_diameter_spfa.test.cpp

@@ -0,0 +1,54 @@
+// @brief Tree Diameter
+#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
+#pragma GCC optimize("Ofast,unroll-loops")
+#include <iostream>
+#include "blazingio/blazingio.min.hpp"
+#include "cp-algo/graph/shortest_path.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+using namespace cp_algo::graph;
+
+template<weighted_undirected_graph_type graph>
+std::tuple<int64_t, node_index, std::vector<edge_index>> tree_diameter(graph const& g) {
+    auto [d1, _] = spfa(g, 0);
+    node_index s = 0;
+    for(auto v: g.nodes()) {
+        if (d1[v] > d1[s]) {
+            s = v;
+        }
+    }
+    auto [d2, pre] = spfa(g, s);
+    node_index t = 0;
+    for(auto v: g.nodes()) {
+        if (d2[v] > d2[t]) {
+            t = v;
+        }
+    }
+    return {d2[t], s, recover_path(pre, s, t)};
+}
+
+void solve() {
+    int n;
+    cin >> n;
+    weighted_graph g(n);
+    g.read_edges(n - 1);
+    auto [d, s, path] = tree_diameter(g);
+    cout << d << ' ' << size(path) + 1 << '\n';
+    cout << s;
+    for(auto e: path) {
+        cout << ' ' << g.edge(e).to;
+    }
+    cout << "\n";
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t = 1;
+    //cin >> t;
+    while(t--) {
+        solve();
+    }
+}