0w1

CFR 767 C. Garland ( Tree DP )

Problem - C - Codeforces

題意:
給一棵節點帶權的樹。問是否存在一種分割方法,使得可以分成三個非空子樹,使得權重總和相同。輸出方案。

資料規模:
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;
}