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);
}