Rerooting(全方位木DP)
(tree/rerooting.hpp)
全方位木DP
テンプレート
// 「T : 根が virtual である根付き木」に対応する情報を管理する
using T = ;
// 空の状態に対応する情報
T leaf = ;
// T 同士をマージ
auto f1 = [&](T x, T y) -> T {
};
// T の根に頂点 c および辺 c-p を追加する (p は virtual)
auto f2 = [&](T x, int c, int p) -> T {
};
Rerooting<T, decltype(g), decltype(f1), decltype(f2)> dp(g, f1, f2, leaf);
Depends on
Verified with
Code
#pragma once
#include "../graph/graph-template.hpp"
// Rerooting
// f1(c1, c2) ... merge value of child node
// f2(memo[i], chd, par) ... return value from child node to parent node
// memo[i] ... result of subtree rooted i
// dp[i] ... result of tree rooted i
template < typename T , typename G , typename F1 , typename F2 >
struct Rerooting {
const G & g ;
const F1 f1 ;
const F2 f2 ;
vector < T > memo , dp ;
T leaf ;
Rerooting ( const G & _g , const F1 _f1 , const F2 _f2 , const T & _leaf )
: g ( _g ), f1 ( _f1 ), f2 ( _f2 ), memo ( g . size ()), dp ( g . size ()), leaf ( _leaf ) {
dfs ( 0 , - 1 );
dfs2 ( 0 , - 1 , T {});
}
const T & operator []( int i ) const { return dp [ i ]; }
void dfs ( int cur , int par ) {
vector < T > chds ;
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
dfs ( dst , cur );
chds . push_back ( f2 ( memo [ dst ], dst , cur ));
}
if ( chds . empty ()) {
memo [ cur ] = leaf ;
} else {
memo [ cur ] = chds [ 0 ];
for ( int i = 1 ; i < ( int ) chds . size (); i ++ ) {
memo [ cur ] = f1 ( memo [ cur ], chds [ i ]);
}
}
}
void dfs2 ( int cur , int par , const T & pval ) {
// get cumulative sum
vector < T > buf ;
if ( cur != 0 ) buf . push_back ( pval );
for ( auto dst : g [ cur ]) {
if ( dst == par ) continue ;
buf . push_back ( f2 ( memo [ dst ], dst , cur ));
}
vector < T > head ( buf . size ()), tail ( buf . size ());
if ( ! buf . empty ()) {
head [ 0 ] = buf [ 0 ];
for ( int i = 1 ; i < ( int ) buf . size (); i ++ ) {
head [ i ] = f1 ( head [ i - 1 ], buf [ i ]);
}
tail . back () = buf . back ();
for ( int i = ( int ) buf . size () - 2 ; i >= 0 ; i -- ) {
tail [ i ] = f1 ( tail [ i + 1 ], buf [ i ]);
}
}
dp [ cur ] = head . empty () ? leaf : head . back ();
int idx = cur == 0 ? 0 : 1 ;
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
T val ;
bool first = idx == 0 ;
bool last = idx + 1 == ( int ) head . size ();
if ( first and last ) {
val = leaf ;
} else if ( first ) {
val = tail [ idx + 1 ];
} else if ( last ) {
val = head [ idx - 1 ];
} else {
val = f1 ( head [ idx - 1 ], tail [ idx + 1 ]);
}
dfs2 ( dst , cur , f2 ( val , cur , dst ));
idx ++ ;
}
}
};
/**
* @brief Rerooting(全方位木DP)
* @docs docs/tree/rerooting.md
*/
#line 2 "tree/rerooting.hpp"
#line 2 "graph/graph-template.hpp"
template < typename T >
struct edge {
int src , to ;
T cost ;
edge ( int _to , T _cost ) : src ( - 1 ), to ( _to ), cost ( _cost ) {}
edge ( int _src , int _to , T _cost ) : src ( _src ), to ( _to ), cost ( _cost ) {}
edge & operator = ( const int & x ) {
to = x ;
return * this ;
}
operator int () const { return to ; }
};
template < typename T >
using Edges = vector < edge < T >> ;
template < typename T >
using WeightedGraph = vector < Edges < T >> ;
using UnweightedGraph = vector < vector < int >> ;
// Input of (Unweighted) Graph
UnweightedGraph graph ( int N , int M = - 1 , bool is_directed = false ,
bool is_1origin = true ) {
UnweightedGraph g ( N );
if ( M == - 1 ) M = N - 1 ;
for ( int _ = 0 ; _ < M ; _ ++ ) {
int x , y ;
cin >> x >> y ;
if ( is_1origin ) x -- , y -- ;
g [ x ]. push_back ( y );
if ( ! is_directed ) g [ y ]. push_back ( x );
}
return g ;
}
// Input of Weighted Graph
template < typename T >
WeightedGraph < T > wgraph ( int N , int M = - 1 , bool is_directed = false ,
bool is_1origin = true ) {
WeightedGraph < T > g ( N );
if ( M == - 1 ) M = N - 1 ;
for ( int _ = 0 ; _ < M ; _ ++ ) {
int x , y ;
cin >> x >> y ;
T c ;
cin >> c ;
if ( is_1origin ) x -- , y -- ;
g [ x ]. emplace_back ( x , y , c );
if ( ! is_directed ) g [ y ]. emplace_back ( y , x , c );
}
return g ;
}
// Input of Edges
template < typename T >
Edges < T > esgraph ( int N , int M , int is_weighted = true , bool is_1origin = true ) {
Edges < T > es ;
for ( int _ = 0 ; _ < M ; _ ++ ) {
int x , y ;
cin >> x >> y ;
T c ;
if ( is_weighted )
cin >> c ;
else
c = 1 ;
if ( is_1origin ) x -- , y -- ;
es . emplace_back ( x , y , c );
}
return es ;
}
// Input of Adjacency Matrix
template < typename T >
vector < vector < T >> adjgraph ( int N , int M , T INF , int is_weighted = true ,
bool is_directed = false , bool is_1origin = true ) {
vector < vector < T >> d ( N , vector < T > ( N , INF ));
for ( int _ = 0 ; _ < M ; _ ++ ) {
int x , y ;
cin >> x >> y ;
T c ;
if ( is_weighted )
cin >> c ;
else
c = 1 ;
if ( is_1origin ) x -- , y -- ;
d [ x ][ y ] = c ;
if ( ! is_directed ) d [ y ][ x ] = c ;
}
return d ;
}
/**
* @brief グラフテンプレート
* @docs docs/graph/graph-template.md
*/
#line 4 "tree/rerooting.hpp"
// Rerooting
// f1(c1, c2) ... merge value of child node
// f2(memo[i], chd, par) ... return value from child node to parent node
// memo[i] ... result of subtree rooted i
// dp[i] ... result of tree rooted i
template < typename T , typename G , typename F1 , typename F2 >
struct Rerooting {
const G & g ;
const F1 f1 ;
const F2 f2 ;
vector < T > memo , dp ;
T leaf ;
Rerooting ( const G & _g , const F1 _f1 , const F2 _f2 , const T & _leaf )
: g ( _g ), f1 ( _f1 ), f2 ( _f2 ), memo ( g . size ()), dp ( g . size ()), leaf ( _leaf ) {
dfs ( 0 , - 1 );
dfs2 ( 0 , - 1 , T {});
}
const T & operator []( int i ) const { return dp [ i ]; }
void dfs ( int cur , int par ) {
vector < T > chds ;
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
dfs ( dst , cur );
chds . push_back ( f2 ( memo [ dst ], dst , cur ));
}
if ( chds . empty ()) {
memo [ cur ] = leaf ;
} else {
memo [ cur ] = chds [ 0 ];
for ( int i = 1 ; i < ( int ) chds . size (); i ++ ) {
memo [ cur ] = f1 ( memo [ cur ], chds [ i ]);
}
}
}
void dfs2 ( int cur , int par , const T & pval ) {
// get cumulative sum
vector < T > buf ;
if ( cur != 0 ) buf . push_back ( pval );
for ( auto dst : g [ cur ]) {
if ( dst == par ) continue ;
buf . push_back ( f2 ( memo [ dst ], dst , cur ));
}
vector < T > head ( buf . size ()), tail ( buf . size ());
if ( ! buf . empty ()) {
head [ 0 ] = buf [ 0 ];
for ( int i = 1 ; i < ( int ) buf . size (); i ++ ) {
head [ i ] = f1 ( head [ i - 1 ], buf [ i ]);
}
tail . back () = buf . back ();
for ( int i = ( int ) buf . size () - 2 ; i >= 0 ; i -- ) {
tail [ i ] = f1 ( tail [ i + 1 ], buf [ i ]);
}
}
dp [ cur ] = head . empty () ? leaf : head . back ();
int idx = cur == 0 ? 0 : 1 ;
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
T val ;
bool first = idx == 0 ;
bool last = idx + 1 == ( int ) head . size ();
if ( first and last ) {
val = leaf ;
} else if ( first ) {
val = tail [ idx + 1 ];
} else if ( last ) {
val = head [ idx - 1 ];
} else {
val = f1 ( head [ idx - 1 ], tail [ idx + 1 ]);
}
dfs2 ( dst , cur , f2 ( val , cur , dst ));
idx ++ ;
}
}
};
/**
* @brief Rerooting(全方位木DP)
* @docs docs/tree/rerooting.md
*/
Back to top page