refactor: dfs once to solve tree_diam

src/code/tree/diam.hpp

@@ -1,21 +1,45 @@
 #ifndef TIFALIBS_TREE_DIAM
 #define TIFALIBS_TREE_DIAM
 
-#include "dfs_info.hpp"
+#include "../graph/dfs.hpp"
 
 namespace tifa_libs::graph {
 
 // {u, v, diam}
 template <class G>
-constexpr auto tree_diam(G &tree) {
-  auto _ = tree.root;
-  tree_dfs_info<G> info;
-  auto d = info.template reset_dfs_info<td_dis>(tree).dis;
-  u32 u = tree.root = (u32)std::distance(d.begin(), std::max_element(d.begin(), d.end()));
-  d = info.template reset_dfs_info<td_dis>(tree).dis;
-  u32 v = (u32)std::distance(d.begin(), std::max_element(d.begin(), d.end()));
-  tree.root = _;
-  return std::make_tuple(u, v, d[v]);
+constexpr auto tree_diam(G const& tree) {
+  using T = typename G::weight_type;
+
+  u32 n = (u32)tree.g.size();
+  vec<T> mdis(n), mdis2(n);
+  vec<u32> midx(n), midx2(n);
+  std::iota(midx.begin(), midx.end(), 0), std::iota(midx2.begin(), midx2.end(), 0);
+
+  auto init = [](u32, u32) {};
+  auto pre = [](u32, u32) {};
+  auto prew = [](u32, T const&, u32) {};
+  auto post = [&](u32 to, u32 u) {
+    if (T _ = mdis[to] + 1; _ > mdis[u]) {
+      mdis2[u] = mdis[u], mdis[u] = _;
+      midx2[u] = midx[u], midx[u] = midx[to];
+    } else if (_ > mdis2[u]) mdis2[u] = _, midx2[u] = midx2[to];
+  };
+  auto postw = [&](u32 to, T const& w, u32 u) {
+    if (T _ = mdis[to] + w; _ > mdis[u]) {
+      mdis2[u] = mdis[u], mdis[u] = _;
+      midx2[u] = midx[u], midx[u] = midx[to];
+    } else if (_ > mdis2[u]) mdis2[u] = _, midx2[u] = midx2[to];
+  };
+  auto ret = [](u32, u32) {};
+
+  if constexpr (std::is_base_of_v<alist, G>) dfs(tree, tree.root, init, pre, post, ret);
+  else dfs(tree, tree.root, init, prew, postw, ret);
+
+  u32 u = midx[0], v = midx2[0];
+  T d = mdis[0] + mdis2[0];
+  for (u32 i = 1; i < n; ++i)
+    if (d < mdis[i] + mdis2[i]) u = midx[i], v = midx2[i], d = mdis[i] + mdis2[i];
+  return std::make_tuple(u, v, d);
 }
 
 }  // namespace tifa_libs::graph