0w1

CFR 486 D. Valid Sets ( DP, Tree )

Problem - D - Codeforces

題意:
有一棵 n 個節點的樹,各自有一個值 a[ i ]。給你一個 d,求有多少邊和點構成的連通塊集合,滿足最大值和最小值的差不超過 d,模上 1e9 + 7。

資料規模:
0 ≤ d ≤ 2000
1 ≤ n ≤ 2000
1 ≤ ai ≤ 2000
TL: 1000 ms
ML: 256 MB

解法:
顯然可以 O( N ^ 2 )。考慮枚舉每個節點,為集合中最小值。如果給原先的問題附加一個條件,給你集合中最小值,感覺可解一點。再考慮指定某個節點必須被集合包含且同時為最小值的限制,這個就很好算:以該節點 i 為根做 dfs,dp[ x ] 代表以 x 為根且包含 x 的子樹有多少連通集合使得每個點不小於 i 的值,且和 i 的值的差不超過 d,這麼一來 dp[ u ] = Pi{ dp[ v ] + 1, for all such v that is child of u }。但問題在於,當枚舉的根的值和其他節點的值有相同,就會有重複計算。因此枚舉根 i 後,當進行 dfs,若即將拜訪到和 i 的值相同的某個節點 j,當且僅當 i < j 才拜訪。

時間 / 空間複雜度:
O( N ^ 2 ) / O( N )

const int MOD7 = ( int ) 1e9 + 7;

int D, N;
vi A;
vvi G;

void init(){
  cin >> D >> N;
  A = vi( N );
  for( int i = 0; i < N; ++i ){
    cin >> A[ i ];
  }
  G = vvi( N );
  for( int i = 0; i < N - 1; ++i ){
    int u, v; cin >> u >> v;
    G[ u - 1 ].emplace_back( v - 1 );
    G[ v - 1 ].emplace_back( u - 1 );
  }
}

int ans = 0;

int dfs( int u, int root, vi &vis ){
  int res = 1;
  for( int v : G[ u ] ){
    if( vis[ v ] ) continue;
    vis[ v ] = 1;
    if( A[ v ] == A[ root ] and v < root ) continue;
    if( A[ v ] < A[ root ] ) continue;
    if( A[ v ] - A[ root ] > D ) continue;
    res = 1LL * res * ( 1 + dfs( v, root, vis ) ) % MOD7;
  }
  return res;
}

void preprocess(){
  for( int i = 0; i < N; ++i ){
    vi vis( N );
    vis[ i ] = 1;
    ( ans += dfs( i, i, vis ) ) %= MOD7;
  }
}

void solve(){
  cout << ans << endl;
}