1 files changed, 41 insertions, 44 deletions

diff --git a/content/graph/reroot.cpp b/content/graph/reroot.cpp
index b3571f1..cd63b3d 100644
--- a/content/graph/reroot.cpp
+++ b/content/graph/reroot.cpp
@@ -1,51 +1,48 @@
-// output[r] = dp[r], where dp[v] :=
-//  fin(Sum_{child c of v, regarding root r} fromChild( dp[c] ))
+using W = ll; // edge weight type
+vector<vector<pair<int, W>>> adj;
 
-// To remove weights, remove every "w" and "W" and fix errors
 struct Reroot {
-  using D = /*todo*/; // dp value
-  using A = /*todo(often D)*/; //value from a vertex's child(ren)
-                               //(A,agg,e) commutative monoid
+	using T = ll; // dp type
 
-  A e = /*todo*/;
-  A fromChild(int v, int c, auto w, D dp_c) { /*todo*/ }
-  static A agg(A a, A b) { /*todo*/ }
-  D fin(int v, A chilsAgg) { /*todo*/ }
+	static constexpr T E = 0; // neutral element
+	T takeChild(int v, int c, W w, T x) {} // move child along edge
+	static T comb(T x, T y) {}
+	T fin(int v, T x) {} // add v to own dp value x
 
-  vector<D> dp;
+	vector<T> dp;
 
-  D dfs0(auto& g, int v, int from = -1) {
-    A ca = e;
-    for (auto [c, w] : g[v]) {
-      if (c == from) continue;
-      ca = agg(ca, fromChild(v, c, w, dfs0(c, v, g)));
-    }
-    return dp[v] = fin(v, ca);
-  }
+	T dfs0(int v, int from = -1) {
+		T val = E;
+		for (auto& [u, w] : adj[v]) {
+			if (u == from) continue;
+			val = comb(val, takeChild(v, u, w, dfs0(u, v)));
+		}
+		return dp[v] = fin(v, val);
+	}
 
-  void dfs1(auto& g, int v, int from = -1) {
-    vector ps = {e};
-    for (auto [c, w] : g[v]) {
-      ps.push_back(fromChild(v, c, w, dp[c]));
-    }
-    auto ss = ps;
-    exclusive_scan(be(ps), ps.begin(), e, agg);
-    exclusive_scan(ss.rbegin(), ss.rend(), ss.rbegin(), e, agg);
-    int i = 0;
-    for (auto [c, w] : g[v]) {
-      ++i;
-      if (c == from) continue;
-      dp[v] = fin(v, agg(ss[i], ps[i]));
-      dfs1(c, v, g);
-    }
-    dp[v] = fin(v, ss[0]);
-  }
+	void dfs1(int v, int from = -1) {
+		vector<T> pref = {E};
+		for (auto [u, w] : adj[v]) {
+			pref.push_back(takeChild(v, u, w, dp[u]));
+		}
+		auto suf = pref;
+		partial_sum(all(pref), pref.begin(), comb);
+		exclusive_scan(suf.rbegin(), suf.rend(),
+		               suf.rbegin(), E, comb);
 
-  template<typename W>
-  auto solve(vector<vector<pair<int, W>>> g) {
-    dp.resize(sz(g));
-    dfs0(g, 0);
-    dfs1(g, 0);
-    return dp;
-  }
-};
\ No newline at end of file
+		for (int i = 0; i < sz(adj[v]); i++) {
+			auto [u, w] = adj[v][i];
+			if (u == from) continue;
+			dp[v] = fin(v, comb(pref[i], suf[i + 1]));
+			dfs1(u, v);
+		}
+		dp[v] = fin(v, suf[0]);
+	}
+
+	auto solve() {
+		dp.assign(sz(adj), E);
+		dfs0(0);
+		dfs1(0);
+		return dp;
+	}
+};