88// Reference:
99// - https://atcoder.jp/contests/abc222/editorial/2749
1010// - https://null-mn.hatenablog.com/entry/2020/04/14/124151
11- template <class Edge , class St , class Ch , Ch (*merge)(Ch, Ch), Ch (*f)(St, int , Edge),
12- St (*g)(Ch, int ), Ch (*e)()>
11+ template <class Edge , class Subtree , class Children , Children (*rake)(Children, Children),
12+ Children (*add_edge)(Subtree, int , Edge), Subtree (*add_vertex)(Children, int ),
13+ Children (*e)()>
1314struct rerooting {
1415 int n_;
1516 std::vector<int > par, visited;
1617 std::vector<std::vector<std::pair<int , Edge>>> to;
17- std::vector<St> dp_subtree;
18- std::vector<St> dp_par;
19- std::vector<St> dpall;
18+
19+ // dp_subtree[i] = DP(root=i, edge (i, par[i]) is removed).
20+ std::vector<Subtree> dp_subtree;
21+
22+ // dp_par[i] = DP(root=par[i], edge (i, par[i]) is removed). dp_par[root] is meaningless.
23+ std::vector<Subtree> dp_par;
24+
25+ // dpall[i] = DP(root=i, all edges exist).
26+ std::vector<Subtree> dpall;
27+
2028 rerooting (const std::vector<std::vector<std::pair<int , Edge>>> &to_)
2129 : n_ (to_.size ()), par (n_, -1 ), visited (n_, 0 ), to (to_) {
22- for (int i = 0 ; i < n_; ++i) dp_subtree.push_back (g (e (), i));
30+ for (int i = 0 ; i < n_; ++i) dp_subtree.push_back (add_vertex (e (), i));
2331 dp_par = dpall = dp_subtree;
2432 }
2533
2634 void run_connected (int root) {
27- if (visited[ root] ) return ;
28- visited[ root] = 1 ;
35+ if (visited. at ( root) ) return ;
36+ visited. at ( root) = 1 ;
2937 std::vector<int > visorder{root};
3038
3139 for (int t = 0 ; t < int (visorder.size ()); ++t) {
32- int now = visorder[t];
33- for (const auto &edge : to[now]) {
34- int nxt = edge.first ;
35- if (visited[nxt]) continue ;
40+ int now = visorder.at (t);
41+ for (const auto &[nxt, _] : to[now]) {
42+ if (visited.at (nxt)) continue ;
3643 visorder.push_back (nxt);
37- visited[ nxt] = 1 ;
38- par[ nxt] = now;
44+ visited. at ( nxt) = 1 ;
45+ par. at ( nxt) = now;
3946 }
4047 }
4148
4249 for (int t = int (visorder.size ()) - 1 ; t >= 0 ; --t) {
43- int now = visorder[t];
44- Ch ch = e ();
45- for (const auto &edge : to[now]) {
46- int nxt = edge.first ;
47- if (nxt == par[now]) continue ;
48- ch = merge (ch, f (dp_subtree[nxt], nxt, edge.second ));
50+ const int now = visorder.at (t);
51+ Children ch = e ();
52+ for (const auto &[nxt, edge] : to.at (now)) {
53+ if (nxt != par.at (now)) ch = rake (ch, add_edge (dp_subtree.at (nxt), nxt, edge));
4954 }
50- dp_subtree[ now] = g (ch, now);
55+ dp_subtree. at ( now) = add_vertex (ch, now);
5156 }
5257
53- std::vector<Ch > left;
58+ std::vector<Children > left;
5459 for (int now : visorder) {
55- int m = int (to[ now] .size () );
60+ const int m = to. at ( now) .size ();
5661 left.assign (m + 1 , e ());
5762 for (int j = 0 ; j < m; j++) {
58- int nxt = to[ now][j]. first ;
59- const St &st = (nxt == par[ now] ? dp_par[ now] : dp_subtree[ nxt] );
60- left[ j + 1 ] = merge (left[j], f (st, nxt, to[now][j]. second ));
63+ const auto &[ nxt, edge] = to. at ( now). at (j) ;
64+ const Subtree &st = (nxt == par. at ( now) ? dp_par. at ( now) : dp_subtree. at ( nxt) );
65+ left. at ( j + 1 ) = rake (left. at (j), add_edge (st, nxt, edge ));
6166 }
62- dpall[ now] = g (left.back (), now);
67+ dpall. at ( now) = add_vertex (left.back (), now);
6368
64- Ch rprod = e ();
69+ Children rprod = e ();
6570 for (int j = m - 1 ; j >= 0 ; --j) {
66- int nxt = to[now][j].first ;
67- if (nxt != par[now]) dp_par[nxt] = g (merge (left[j], rprod), now);
71+ const auto &[nxt, edge] = to.at (now).at (j);
72+
73+ if (nxt != par.at (now)) dp_par.at (nxt) = add_vertex (rake (left.at (j), rprod), now);
6874
69- const St &st = (nxt == par[ now] ? dp_par[ now] : dp_subtree[ nxt] );
70- rprod = merge ( f (st, nxt, to[now][j]. second ), rprod);
75+ const Subtree &st = (nxt == par. at ( now) ? dp_par. at ( now) : dp_subtree. at ( nxt) );
76+ rprod = rake ( add_edge (st, nxt, edge ), rprod);
7177 }
7278 }
7379 }
7480
7581 void run () {
7682 for (int i = 0 ; i < n_; ++i) {
77- if (!visited[i] ) run_connected (i);
83+ if (!visited. at (i) ) run_connected (i);
7884 }
7985 }
8086
81- const St &get_subtree (int root_, int par_) const {
87+ const Subtree &get_subtree (int root_, int par_) const {
8288 if (par_ < 0 ) return dpall.at (root_);
8389 if (par.at (root_) == par_) return dp_subtree.at (root_);
8490 if (par.at (par_) == root_) return dp_par.at (par_);
@@ -87,13 +93,13 @@ struct rerooting {
8793};
8894/* Template:
8995struct Subtree {};
90- struct Child {};
96+ struct Children {};
9197struct Edge {};
92- Child e() { return Child (); }
93- Child merge(Child x, Child y) { return Child (); }
94- Child f (Subtree x, int ch_id, Edge edge) { return Child (); }
95- Subtree g(Child x, int v_id) { return Subtree(); }
98+ Children e() { return Children (); }
99+ Children rake(Children x, Children y) { return Children (); }
100+ Children add_edge (Subtree x, int ch_id, Edge edge) { return Children (); }
101+ Subtree add_vertex(Children x, int v_id) { return Subtree(); }
96102
97103vector<vector<pair<int, Edge>>> to;
98- rerooting<Edge, Subtree, Child, merge, f, g , e> tree(to);
104+ rerooting<Edge, Subtree, Children, rake, add_edge, add_vertex , e> tree(to);
99105*/
0 commit comments