Contact Form Submission - Suggestion (Gold - DP on Trees - Solving For All Roots)

[maggieliu05]

Someone submitted the contact form!

URL: https://usaco.guide/gold/all-roots?lang=cpp
Module: Gold - DP on Trees - Solving For All Roots
Topic: Suggestion
Message:
Worth mentioning rerooting templates that allow implementing such problems more quickly and reliably in linear time. Example: https://codeforces.com/blog/entry/124286 (though this is NlogN rather than N)

[bqi343]

Example: solving Tree Distances I with a template modified from the blog to run in linear time.

• Only the code in main is problem-specific.
• The module solution keeps track of both first and second maximums for each subtree, because it needs to account for how the longest path changes after removing a subtree from another. This is easy to mess up if you're not being careful. On the other hand, the code below only requires logic for merging subtrees sharing a common root to be provided, which is trivial (just take the longer path).

...

namespace reroot {

auto rerooter = [](const auto &g, const auto &init_a, const auto &v_to_a,
                   const auto &a_to_v, const auto &add) {
	using Agg = decay_t<decltype(init_a())>;
	using Val = decay_t<decltype(a_to_v(init_a(), 0))>;
	int N = sz(g);
	vi par(N);
	V<Val> dp(N), dp_root(N);
	V<V<Val>> dp_down(N), dp_up(N);
	y_combinator([&](auto dfs_down, int x) -> void {
		Agg agg = init_a();
		F0R(e, sz(g[x])) {
			int y = g[x][e];
			if (y != par[x]) {
				par[y] = x;
				dfs_down(y);
				agg = add(agg, v_to_a(dp[y], x, e));
			}
		}
		dp[x] = a_to_v(agg, x);
	})(0);
	y_combinator([&](auto dfs_up, int x) -> void {
		dp[par[x]] = dp[x];
		V<Agg> pref_aggs{init_a()}, suf_aggs{init_a()};
		F0R(e, sz(g[x])) {
			int y = g[x][e];
			pref_aggs.pb(add(pref_aggs.bk, v_to_a(dp[y], x, e)));
		}
		R0F(e, sz(g[x])) {
			int y = g[x][e];
			suf_aggs.pb(add(suf_aggs.bk, v_to_a(dp[y], x, e)));
		}
		dp_root[x] = a_to_v(pref_aggs.bk, x);
		F0R(e, sz(g[x])) {
			int y = g[x][e];
			dp_down[x].pb(dp[y]);
			dp_up[x].pb(
			    a_to_v(add(pref_aggs[e], suf_aggs[sz(g[x]) - 1 - e]), x));
			if (y != par[x]) {
				dp[y] = dp_up[x][e];
				dfs_up(y);
			}
		}
	})(0);
	return make_tuple(dp_root, dp_down, dp_up);
};

}

int main() {
	setIO();
	def(int, N);
	V<vi> g(N);
	rep(N - 1) {
		def(int, a, b);
		--a, --b;
		g[a].pb(b), g[b].pb(a);
	}
	struct Val {  // max depth for subtree including root
		int mx;
	};
	struct Agg {  // max depth for subtree excluding root
		int mx;
	};
	auto [dp_root, dp_down, dp_up] = reroot::rerooter(
	    g,
	    []() -> Agg {  // init empty subtree
		    return {};
	    },
	    [](Val v, int x, int e) -> Agg {  // add edge to subtree
		    return {v.mx};
	    },
	    [](Agg a, int x) -> Val {  // add root to subtree
		    return {a.mx + 1};
	    },
	    [](Agg a, Agg b) -> Agg {  // merge subtrees
		    return {max(a.mx, b.mx)};
	    });
	vi ans;
	F0R(i, N) ans.pb(dp_root.at(i).mx - 1);
	ps(ans);
}