[Add] rerooting-dp.hpp

cpp/rerooting-dp.hpp

@@ -0,0 +1,100 @@
+#pragma once
+
+/**
+ * @file rerooting-dp.hpp
+ * @brief 全方位木DP
+ */
+#include <algorithm>
+#include <type_traits>
+
+#include "tree.hpp"
+
+/**
+ * @brief 全方位木DP
+ * dp[u] = addnode(merge(addedge(dp[v1], g_uv1, uv1.id), merge(addedge(dp[v2], g_uv2, uv2.id), ...)), u) v_iはuの子
+ *
+ * @tparam E 可換モノイド
+ * @tparam V DPの型
+ * @param tree
+ * @param e 可換モノイドの単位元
+ * @param merge (E,E)->E 可換
+ * @param addedge (V,Cost,int)->E
+ * @param addnode (E,int)->V
+ * @return std::vector<V>
+ */
+template <typename E, typename V, class Merge, class AddEdge, class AddNode, typename Cost>
+std::vector<V> rerooting_dp(const Tree<Cost>& tree, E e, Merge merge, AddEdge addedge, AddNode addnode) {
+ static_assert(std::is_invocable_r_v<E, Merge, E, E>);
+ static_assert(std::is_invocable_r_v<E, AddEdge, V, Cost, int>);
+ static_assert(std::is_invocable_r_v<V, AddNode, E, int>);
+ assert(tree.n > 0);
+ RootedTree<Cost> rooted(tree, 0);
+ std::vector<V> subdp(rooted.n);
+ for (int u : rooted.postorder) {
+ E tmp = e;
+ for (auto edge : rooted.child[u]) {
+ tmp = merge(tmp, addedge(subdp[edge.dst], edge.cost, edge.id));
+ }
+ subdp[u] = addnode(tmp, u);
+ }
+ std::vector<V> dp(rooted.n);
+ std::vector<E> pe(rooted.n, e);
+ for (int u : rooted.preorder) {
+ const auto& ch = rooted.child[u];
+ std::vector<E> ri(ch.size() + 1, e);
+ for (size_t i = ch.size(); i > 0; i--) {
+ ri[i - 1] = merge(ri[i], addedge(subdp[ch[i - 1].dst], ch[i - 1].cost, ch[i - 1].id));
+ }
+ dp[u] = addnode(merge(pe[u], ri[0]), u);
+ E le = e;
+ for (size_t i = 0; i < ch.size(); i++) {
+ pe[ch[i].dst] = addedge(addnode(merge(pe[u], merge(le, ri[i + 1])), u), ch[i].cost, ch[i].id);
+ le = merge(le, addedge(subdp[ch[i].dst], ch[i].cost, ch[i].id));
+ }
+ }
+ return dp;
+}
+
+/**
+ * @brief 最も遠い頂点までの距離
+ */
+template <typename Cost>
+inline std::vector<Cost> farthest_node_dist(const Tree<Cost>& tree) {
+ return rerooting_dp<Cost, Cost>(
+ tree,
+ Cost{},
+ [](Cost a, Cost b) { return std::max<Cost>(a, b); },
+ [](Cost a, Cost c, int i) { return a + c; },
+ [](Cost a, int i) { return a; });
+}
+
+/**
+ * @brief 他の全ての頂点との距離の合計
+ */
+template <typename Cost>
+inline std::vector<Cost> distance_sums(const Tree<Cost>& tree) {
+ using P = typename std::pair<Cost, int>;
+ auto tmp = rerooting_dp<P, P>(
+ tree,
+ P{{}, 0},
+ [](P a, P b) { return P{a.first + b.first, a.second + b.second}; },
+ [](P a, Cost c, int i) { return P{a.first + a.second * c, a.second}; },
+ [](P a, int i) { return P{a.first, a.second + 1}; });
+ std::vector<Cost> res;
+ res.reserve(tree.n);
+ for (const P& s : tmp) res.push_back(s.first);
+ return res;
+}
+
+/**
+ * @brief その頂点を含む連結な部分グラフの個数
+ */
+template <typename V, typename Cost>
+inline std::vector<V> connected_subgraph_count(const Tree<Cost>& tree) {
+ return rerooting_dp<V, V>(
+ tree,
+ V(1),
+ [](V a, V b) { return a * b; },
+ [](V a, Cost c, int i) { return a + 1; },
+ [](V a, int i) { return a; });
+}

test/aoj-1595.test.cpp

@@ -0,0 +1,16 @@
+#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1595"
+
+#include <iostream>
+
+#include "../cpp/rerooting-dp.hpp"
+
+int main() {
+ int N, M;
+ std::cin >> N;
+ Tree<long long> tree(N);
+ tree.read();
+ auto ans = farthest_node_dist(tree);
+ for (int i = 0; i < N; i++) {
+ std::cout << 2 * N - 2 - ans[i] << std::endl;
+ }
+}