From 8bad460ff55d67bac81ae31196b05c1524a52e34 Mon Sep 17 00:00:00 2001 From: f1or1an <110727007+f1or1an@users.noreply.github.com> Date: Thu, 29 Aug 2024 14:29:10 +0000 Subject: shorter rerooting (tested on https://codeforces.com/gym/532576/submission/278652360) --- content/graph/reroot.cpp | 99 ++++++++++++++++++++---------------------------- 1 file changed, 42 insertions(+), 57 deletions(-) diff --git a/content/graph/reroot.cpp b/content/graph/reroot.cpp index 4c6a748..a413917 100644 --- a/content/graph/reroot.cpp +++ b/content/graph/reroot.cpp @@ -1,62 +1,47 @@ -// Usual Tree DP can be broken down in 4 steps: -// - Initialize dp[v] = identity -// - Iterate over all children w and take a value for w -// by looking at dp[w] and possibly the edge label of v -> w -// - combine the values of those children -// usually this operation should be commutative and associative -// - finalize the dp[v] after iterating over all children +// input: undirected (un)weighted tree as +// adjacency list containing pairs +// (To remove weights, remove every "w" and fix errors) +// output[r] = dp[r], where dp[v] := +// fin(Sum_{child c of v, regarding root r} from_child( dp[c] )) struct Reroot { - using T = ll; + using D = todo; // dp value + using A = todo (often D); // value from a vertex's child(ren) + // (A,agg,e) commutative monoid - // identity element - T E() {} - // x: dp value of child - // e: index of edge going to child - T takeChild(T x, int e) {} - T comb(T x, T y) {} - // called after combining all dp values of children - T fin(T x, int v) {} + A e = todo; + A from_child(z v, z c, auto w, D dp_c) { todo } + static A agg(A a, A b) { todo } + D fin(z v, A chils_agg) { todo } - vector>> g; - vector ord, pae; - vector dp; + vector dp; - T dfs(int v) { - ord.push_back(v); - for (auto [w, e] : g[v]) { - g[w].erase(find(all(g[w]), pair(v, e^1))); - pae[w] = e^1; - dp[v] = comb(dp[v], takeChild(dfs(w), e)); - } - return dp[v] = fin(dp[v], v); - } + D dfs0(z v, z p, auto& g) { + A ca = e; + for(auto [c, w] : g[v]) if(c-p) { + ca = agg(ca, from_child(v, c, w, dfs0(c, v, g))); + } + return dp[v] = fin(v, ca); + } + void dfs1(z v, z p, auto& g) { + vector ps = {e}; + for(auto [c, w] : g[v]) { + ps.push_back(from_child(v, c, w, dp[c])); + } + auto ss = ps; + exclusive_scan(ps.begin(), ps.end(), ps.begin(), e, agg); + exclusive_scan(ss.rbegin(),ss.rend(),ss.rbegin(),e, agg); + z i = 0; + for(auto [c, w] : g[v]) if(++i, c-p) { + dp[v] = fin(v, agg(ss[i], ps[i])); + dfs1(c, v, g); + } + dp[v] = fin(v, s[0]); + } - vector solve(int n, vector> edges) { - g.resize(n); - for (int i = 0; i < n-1; i++) { - g[edges[i].first].emplace_back(edges[i].second, 2*i); - g[edges[i].second].emplace_back(edges[i].first, 2*i+1); - } - pae.assign(n, -1); - dp.assign(n, E()); - dfs(0); - vector updp(n, E()), res(n, E()); - for (int v : ord) { - vector pref(sz(g[v])+1), suff(sz(g[v])+1); - if (v != 0) pref[0] = takeChild(updp[v], pae[v]); - for (int i = 0; i < sz(g[v]); i++){ - auto [u, w] = g[v][i]; - pref[i+1] = suff[i] = takeChild(dp[u], w); - pref[i+1] = comb(pref[i], pref[i+1]); - } - for (int i = sz(g[v])-1; i >= 0; i--) { - suff[i] = comb(suff[i], suff[i+1]); - } - for (int i = 0; i < sz(g[v]); i++) { - updp[g[v][i].first] = fin(comb(pref[i], suff[i+1]), v); - } - res[v] = fin(pref.back(), v); - } - return res; - } -}; + auto solve(auto g) { + dp.resize(sz(g)); + dfs0(0, 0, g); + dfs1(0, 0, g); + return dp; + } +}; \ No newline at end of file -- cgit v1.2.3 From 9cb227598ad234265b44ed5feb9890f5401f8e81 Mon Sep 17 00:00:00 2001 From: florian Date: Thu, 29 Aug 2024 16:45:22 +0200 Subject: spacing --- content/graph/reroot.cpp | 68 ++++++++++++++++++++++++------------------------ 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/content/graph/reroot.cpp b/content/graph/reroot.cpp index a413917..2ec9ad4 100644 --- a/content/graph/reroot.cpp +++ b/content/graph/reroot.cpp @@ -4,44 +4,44 @@ // output[r] = dp[r], where dp[v] := // fin(Sum_{child c of v, regarding root r} from_child( dp[c] )) 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 D = todo; // dp value + using A = todo (often D); // value from a vertex's child(ren) + // (A,agg,e) commutative monoid - A e = todo; - A from_child(z v, z c, auto w, D dp_c) { todo } - static A agg(A a, A b) { todo } - D fin(z v, A chils_agg) { todo } + A e = todo; + A from_child(z v, z c, auto w, D dp_c) { todo } + static A agg(A a, A b) { todo } + D fin(z v, A chils_agg) { todo } - vector dp; + vector dp; - D dfs0(z v, z p, auto& g) { - A ca = e; - for(auto [c, w] : g[v]) if(c-p) { - ca = agg(ca, from_child(v, c, w, dfs0(c, v, g))); - } - return dp[v] = fin(v, ca); + D dfs0(z v, z p, auto& g) { + A ca = e; + for (auto [c, w] : g[v]) if(c-p) { + ca = agg(ca, from_child(v, c, w, dfs0(c, v, g))); } - void dfs1(z v, z p, auto& g) { - vector ps = {e}; - for(auto [c, w] : g[v]) { - ps.push_back(from_child(v, c, w, dp[c])); - } - auto ss = ps; - exclusive_scan(ps.begin(), ps.end(), ps.begin(), e, agg); - exclusive_scan(ss.rbegin(),ss.rend(),ss.rbegin(),e, agg); - z i = 0; - for(auto [c, w] : g[v]) if(++i, c-p) { - dp[v] = fin(v, agg(ss[i], ps[i])); - dfs1(c, v, g); - } - dp[v] = fin(v, s[0]); + return dp[v] = fin(v, ca); + } + void dfs1(z v, z p, auto& g) { + vector ps = {e}; + for (auto [c, w] : g[v]) { + ps.push_back(from_child(v, c, w, dp[c])); } - - auto solve(auto g) { - dp.resize(sz(g)); - dfs0(0, 0, g); - dfs1(0, 0, g); - return dp; + auto ss = ps; + exclusive_scan(ps.begin(), ps.end(), ps.begin(), e, agg); + exclusive_scan(ss.rbegin(),ss.rend(),ss.rbegin(),e, agg); + z i = 0; + for (auto [c, w] : g[v]) if(++i, c-p) { + dp[v] = fin(v, agg(ss[i], ps[i])); + dfs1(c, v, g); } + dp[v] = fin(v, s[0]); + } + + auto solve(auto g) { + dp.resize(sz(g)); + dfs0(0, 0, g); + dfs1(0, 0, g); + return dp; + } }; \ No newline at end of file -- cgit v1.2.3 From 36a975670f60d17e80c32ea7bccc0040680dd9c5 Mon Sep 17 00:00:00 2001 From: florian Date: Thu, 29 Aug 2024 16:56:45 +0200 Subject: fix rerooting --- content/graph/reroot.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/content/graph/reroot.cpp b/content/graph/reroot.cpp index 2ec9ad4..59fea94 100644 --- a/content/graph/reroot.cpp +++ b/content/graph/reroot.cpp @@ -35,7 +35,7 @@ struct Reroot { dp[v] = fin(v, agg(ss[i], ps[i])); dfs1(c, v, g); } - dp[v] = fin(v, s[0]); + dp[v] = fin(v, ss[0]); } auto solve(auto g) { -- cgit v1.2.3