CFR 767 C. Garland ( Tree DP )
題意:
給一棵節點帶權的樹。問是否存在一種分割方法,使得可以分成三個非空子樹,使得權重總和相同。輸出方案。
資料規模:
The first line contains single integer n (3 ≤ n ≤ 1e6) — the number of lamps in the garland.
Then n lines follow. The i-th of them contain the information about the i-th lamp: the number lamp ai, it is hanging on (and 0, if is there is no such lamp), and its temperature ti ( - 100 ≤ ti ≤ 100). The lamps are numbered from 1 to n.
解法:
bsum[ i ] : 以 i 為根的子樹的權重總和
ok[ i ] : 以 i 為根的子樹中是否存在一個子樹,其總和為原樹權重總和的三分之一?有的話為該子樹對應的根節點,有多個時優先取深度大
接著分開討論兩種可能的情形:
1. 割處在兩顆不同的子樹:
那麼只要每次判斷是否有兩個以上的子樹 ok
2. 割處被包含在另一個割處的子樹內:
若存在一個割處 v,其權重總和為原樹權重總和的三分之二,而且其 ok[ v ] ≠ v,那麼 v 和 ok[ v ] 為解。
時間 / 空間複雜度:
O( N )
#include <bits/stdc++.h> using namespace std; const int MAXN = ( int ) 1e6; int N; vector< int > G[ MAXN ]; int B[ MAXN ]; int root; int btot; int bsum[ MAXN ]; int ok[ MAXN ]; void dfs_bsum( int u, int fa ){ bsum[ u ] = B[ u ]; for( int v : G[ u ] ){ if( v == fa ) continue; dfs_bsum( v, u ); bsum[ u ] += bsum[ v ]; if( ok[ v ] != -1 ){ ok[ u ] = ok[ v ]; } } if( ok[ u ] == -1 and bsum[ u ] * 3 == btot ){ ok[ u ] = u; } } void dfs( int u, int fa, int up ){ vector< int > ans; for( int v : G[ u ] ){ if( v == fa ) continue; dfs( v, u, up + bsum[ u ] - B[ u ] - bsum[ v ] ); if( ok[ v ] != - 1 ){ if( ok[ v ] != v and bsum[ v ] * 3 == btot * 2 ){ cout << v + 1 << " " << ok[ v ] + 1 << endl; exit( 0 ); } ans.emplace_back( v ); } } if( ans.size() >= 2 ){ cout << ok[ ans[ 0 ] ] + 1 << " " << ok[ ans[ 1 ] ] + 1 << endl; exit( 0 ); } } signed main(){ ios::sync_with_stdio( 0 ); { cin >> N; for( int i = 0; i < N; ++i ){ int fa, t; cin >> fa >> t; B[ i ] = t; btot += t; if( fa ){ G[ fa - 1 ].emplace_back( i ); G[ i ].emplace_back( fa - 1 ); } else{ root = i; } } } { memset( ok, -1, sizeof( ok ) ); dfs_bsum( root, -1 ); dfs( root, -1, B[ root ] ); cout << -1 << endl; } return 0; }