Rerooting(全方位木DP)
(tree/rerooting.hpp)
全方位木DP
テンプレート
using T = ;
// identify element of f1, and answer of leaf
T I = ;
// merge value of child node
auto f1 = [&](T x, T y) -> T {
};
// return value from child node to parent node
auto f2 = [&](T x, int chd, int par) -> T {
};
Rerooting<T, decltype(g), decltype(f1), decltype(f2)> dp(g, f1, f2, I);
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 I ;
Rerooting ( const G & _g , const F1 _f1 , const F2 _f2 , const T & I_ )
: g ( _g ), f1 ( _f1 ), f2 ( _f2 ), memo ( g . size (), I_ ), dp ( g . size (), I_ ), I ( I_ ) {
dfs ( 0 , - 1 );
efs ( 0 , - 1 , I );
}
const T & operator []( int i ) const { return dp [ i ]; }
void dfs ( int cur , int par ) {
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
dfs ( dst , cur );
memo [ cur ] = f1 ( memo [ cur ], f2 ( memo [ dst ], dst , cur ));
}
}
void efs ( int cur , int par , const T & pval ) {
// get cumulative sum
vector < T > buf ;
for ( auto dst : g [ cur ]) {
if ( dst == par ) continue ;
buf . push_back ( f2 ( memo [ dst ], dst , cur ));
}
vector < T > head ( buf . size () + 1 ), tail ( buf . size () + 1 );
head [ 0 ] = tail [ buf . size ()] = I ;
for ( int i = 0 ; i < ( int ) buf . size (); i ++ ) head [ i + 1 ] = f1 ( head [ i ], buf [ i ]);
for ( int i = ( int ) buf . size () - 1 ; i >= 0 ; i -- )
tail [ i ] = f1 ( tail [ i + 1 ], buf [ i ]);
// update
dp [ cur ] = par == - 1 ? head . back () : f1 ( pval , head . back ());
// propagate
int idx = 0 ;
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
efs ( dst , cur , f2 ( f1 ( pval , f1 ( head [ idx ], tail [ idx + 1 ])), 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 I ;
Rerooting ( const G & _g , const F1 _f1 , const F2 _f2 , const T & I_ )
: g ( _g ), f1 ( _f1 ), f2 ( _f2 ), memo ( g . size (), I_ ), dp ( g . size (), I_ ), I ( I_ ) {
dfs ( 0 , - 1 );
efs ( 0 , - 1 , I );
}
const T & operator []( int i ) const { return dp [ i ]; }
void dfs ( int cur , int par ) {
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
dfs ( dst , cur );
memo [ cur ] = f1 ( memo [ cur ], f2 ( memo [ dst ], dst , cur ));
}
}
void efs ( int cur , int par , const T & pval ) {
// get cumulative sum
vector < T > buf ;
for ( auto dst : g [ cur ]) {
if ( dst == par ) continue ;
buf . push_back ( f2 ( memo [ dst ], dst , cur ));
}
vector < T > head ( buf . size () + 1 ), tail ( buf . size () + 1 );
head [ 0 ] = tail [ buf . size ()] = I ;
for ( int i = 0 ; i < ( int ) buf . size (); i ++ ) head [ i + 1 ] = f1 ( head [ i ], buf [ i ]);
for ( int i = ( int ) buf . size () - 1 ; i >= 0 ; i -- )
tail [ i ] = f1 ( tail [ i + 1 ], buf [ i ]);
// update
dp [ cur ] = par == - 1 ? head . back () : f1 ( pval , head . back ());
// propagate
int idx = 0 ;
for ( auto & dst : g [ cur ]) {
if ( dst == par ) continue ;
efs ( dst , cur , f2 ( f1 ( pval , f1 ( head [ idx ], tail [ idx + 1 ])), cur , dst ));
idx ++ ;
}
}
};
/**
* @brief Rerooting(全方位木DP)
* @docs docs/tree/rerooting.md
*/
Back to top page