minimum diameter spanning tree test

adamant-pwn · adamant-pwn · commit c6a9e445f18c · 2025-12-05T18:12:09.000+01:00
diff --git a/cp-algo/graph/ascending_dfs.hpp b/cp-algo/graph/ascending_dfs.hpp
@@ -2,19 +2,15 @@
 #define CP_ALGO_GRAPH_ASCENDING_DFS_HPP
 
 #include "base.hpp"
+#include <cassert>
 #include <vector>
 #include <ranges>
 
 namespace cp_algo::graph {
     // Generic ascending DFS that repeatedly peels leaves (degree 1).
     // `next(v)` must return the unique incident edge of `v` when degree[v] == 1.
     template<undirected_graph_type graph>
-    void ascending_dfs(graph const& tree, auto &&next, auto &&callback, node_index root = 0) {
-        std::vector<int> degree(tree.n());
-        for (auto v : tree.nodes()) {
-            degree[v] = (int)std::ranges::distance(tree.outgoing(v));
-        }
-        degree[root] = 0;
+    void ascending_dfs(graph const& tree, auto &degree, auto &&next, auto &&callback, node_index root) {
         for (auto v : tree.nodes()) {
             while (degree[v] == 1) {
                 edge_index ep = next(v);
@@ -30,13 +26,18 @@ namespace cp_algo::graph {
     template<undirected_graph_type graph>
     auto xor_dfs(graph const& tree, auto &&callback, node_index root = 0) {
         std::vector<edge_index> neig_xor(tree.n());
+        std::vector<int> degree(tree.n());
+        for (auto v : tree.nodes()) {
+            degree[v] = (int)std::ranges::distance(tree.outgoing(v));
+        }
+        degree[root] = 0;
         for (auto v : tree.nodes()) {
             for (auto e : tree.outgoing(v)) {
                 neig_xor[v] ^= e;
             }
         }
         neig_xor[root] ^= edge_index(-1);
-        ascending_dfs(tree, [&](auto v) {
+        ascending_dfs(tree, degree, [&](auto v) {
             edge_index ep = neig_xor[v];
             neig_xor[tree.edge(ep).traverse(v)] ^= ep;
             return ep;
@@ -46,8 +47,20 @@ namespace cp_algo::graph {
     }
     // DFS that uses a precomputed parent-edge array.
     template<undirected_graph_type graph>
-    void parent_dfs(graph const& tree, std::vector<edge_index> const& parent, auto &&callback, node_index root = 0) {
-        ascending_dfs(tree, [&](auto v) {
+    void parent_dfs(graph const& tree, std::vector<edge_index> const& parent, auto &&callback) {
+        std::vector<int> degree(tree.n());
+        node_index root = -1;
+        for (auto [v, e] : parent | std::views::enumerate) {
+            if (e != -1) {
+                degree[v]++;
+                degree[tree.edge(e).traverse(node_index(v))]++;
+            } else {
+                root = node_index(v);
+            }
+        }
+        assert(root != -1);
+        degree[root] = 0;
+        ascending_dfs(tree, degree, [&](auto v) {
             return parent[v];
         }, callback, root);
     }
diff --git a/cp-algo/graph/shortest_path.hpp b/cp-algo/graph/shortest_path.hpp
@@ -105,6 +105,7 @@ namespace cp_algo::graph {
     Context sssp_impl(graph const& g, node_index s) {
         Context context(g.n());
         context.dist[s] = 0;
+        context.pre[s] = -1;
         context.push(s, -1, s);
         while(auto ov = context.next_node()) {
             node_index v = *ov;
diff --git a/cp-algo/graph/tree_diameter.hpp b/cp-algo/graph/tree_diameter.hpp
@@ -5,12 +5,12 @@
 #include <tuple>
 #include <string>
 #include <algorithm>
-#include <memory>
 
 namespace cp_algo::graph {
-    enum diameter_mode { recover_path, no_recover };
-    template<diameter_mode mode = no_recover, weighted_undirected_graph_type graph>
-    auto tree_diameter(graph const& g, std::shared_ptr<const std::vector<edge_index>> parents = nullptr) {
+    enum class diameter_mode { recover_path, no_recover };
+    
+    template<diameter_mode mode = diameter_mode::no_recover, weighted_undirected_graph_type graph>
+    auto tree_diameter(graph const& g, const std::vector<edge_index>* parents = nullptr) {
         struct up_path {
             int64_t length = 0;
             node_index start = 0;
@@ -32,28 +32,31 @@ namespace cp_algo::graph {
                 up[u] = up[v];
             }
         };
+        std::vector<edge_index> parents_owned;
         if (parents) {
             parent_dfs(g, *parents, callback);
         } else {
-            parents = std::make_shared<std::vector<edge_index>>(xor_dfs(g, callback));
+            parents_owned = xor_dfs(g, callback);
+            parents = &parents_owned;
         }
-        if constexpr (mode == no_recover) {
+        if constexpr (mode == diameter_mode::no_recover) {
             return s.length + t.length;
+        } else {
+            auto collect = [&](up_path v) {
+                std::basic_string<edge_index> path;
+                while(v.length) {
+                    edge_index ep = (*parents)[v.start];
+                    path.push_back(ep);
+                    v.length -= g.edge(ep).w;
+                    v.start = g.edge(ep).traverse(v.start);
+                }
+                return path;
+            };
+            auto paths = collect(s);
+            auto patht = collect(t);
+            std::ranges::reverse(patht);
+            return std::tuple{s.length + t.length, s.start, paths += patht};
         }
-        auto collect = [&](up_path v) {
-            std::basic_string<edge_index> path;
-            while(v.length) {
-                edge_index ep = (*parents)[v.start];
-                path.push_back(ep);
-                v.length -= g.edge(ep).w;
-                v.start = g.edge(ep).traverse(v.start);
-            }
-            return path;
-        };
-        auto paths = collect(s);
-        auto patht = collect(t);
-        std::ranges::reverse(patht);
-        return std::tuple{s.length + t.length, s.start, paths += patht};
     }
 }
 #endif // CP_ALGO_GRAPH_TREE_DIAMETER_HPP
diff --git a/verify/graph/minimum_diameter_spanning_tree.test.cpp b/verify/graph/minimum_diameter_spanning_tree.test.cpp
@@ -0,0 +1,52 @@
+// @brief Minimum Diameter Spanning Tree
+#define PROBLEM "https://judge.yosupo.jp/problem/minimum_diameter_spanning_tree"
+#pragma GCC optimize("Ofast,unroll-loops")
+#include <iostream>
+#include "blazingio/blazingio.min.hpp"
+#include "cp-algo/graph/shortest_path.hpp"
+#include "cp-algo/graph/tree_diameter.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+using namespace cp_algo::graph;
+
+template<weighted_undirected_graph_type graph>
+auto minimum_diameter_spanning_tree(graph const& g) {
+    int64_t min_diameter = shortest_path_context::inf;
+    std::vector<edge_index> best_edges;
+    for(auto v: g.nodes()) {
+        auto [_, p] = dijkstra(g, v);
+        auto d = tree_diameter(g, &p);
+        if (d < min_diameter) {
+            min_diameter = d;
+            best_edges = p;
+        }
+    }
+    return std::pair{min_diameter, best_edges};
+}
+
+void solve() {
+    int n, m;
+    cin >> n >> m;
+    weighted_graph g(n);
+    g.read_edges(m);
+    auto [X, E] = minimum_diameter_spanning_tree(g);
+    cout << X << '\n';
+    for(auto e: E) {
+        if (e != -1) {
+            cout << e << " ";
+        }
+    }
+    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();
+    }
+}
diff --git a/verify/graph/tree_diameter.test.cpp b/verify/graph/tree_diameter.test.cpp
@@ -14,7 +14,7 @@ void solve() {
     cin >> n;
     weighted_graph g(n);
     g.read_edges(n - 1);
-    auto [d, v, path] = tree_diameter<recover_path>(g);
+    auto [d, v, path] = tree_diameter<diameter_mode::recover_path>(g);
     cout << d << ' ' << size(path) + 1 << '\n';
     cout << v;
     for(auto e: path) {
