Codeforces 486D Valid Sets：Tree dp【n遍O(n)的dp】

题目链接：http://codeforces.com/problemset/problem/486/D

题意：

　　给你一棵树，n个节点，每个节点的点权为a[i]。

　　问你有多少个连通子图，使得子图中的max(a[i]) - min(a[i]) <= d。

　　ps.连通子图的定义：

　　　　如果一个点集V为一个连通子图，则对于任意两点a,b∈V，有a到b路径上的所有点u∈V。

 

题解：

　　因为要保证max(a[i]) - min(a[i]) <= d，所以可以人为地选出一个点rt作为点权最大的点。

　　这样在求以rt为最大点的连通子图个数时，只用考虑点权不超过a[rt]的点。

 

　　然而如果有两个点i,j的点权相同，分别以i,j作为最大点的两堆子图中会有重复。

　　所以可以定义一下：当以rt作为最大点时，所有点权与a[rt]相等的点，它们的节点编号id[i]必须大于id[rt]。

　　这样就能避免重复了。

 

　　所以最终答案 = ∑（以i为最大点的连通子图个数）

　　求以i为最大点的连通子图个数，只需一遍O(n)的dp就行。

 

　　注意，当前这是一棵无根树。以下所说的“i的子树”意思是：从i出发往叶子方向的那一堆点。

 

　　假设当前以rt作为最大点。　　

　　则加入连通子图的点i必须满足：

　　　　（1）a[rt]-a[i]<=d（保证满足题目条件）

　　　　（2）a[i]<=a[rt]（保证a[rt]为最大点）

　　　　（3）如果a[i]==a[rt] && i!=rt，则要满足rt<i（避免重复计数）

 

　　表示状态：

　　　　dp[i] = numbers

　　　　表示节点i肯定要选，i的子树所构成的合法连通子图个数

　　找出答案：

　　　　每次ans += dp[rt]

　　如何转移：

　　　　dp[i] = ∏ (dp[son]+1)

　　边界条件：

　　　　对于叶子结点leaf: dp[leaf] = 1

 

　　总复杂度O(n^2)。

 

AC Code:

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <vector>
#define MAX_N 2005
#define MOD 1000000007

using namespace std;

int n,d;
int a[MAX_N];
vector<int> edge[MAX_N];

void read()
{
    cin>>d>>n;
    for(int i=1;i<=n;i++) cin>>a[i];
    int x,y;
    for(int i=1;i<n;i++)
    {
        cin>>x>>y;
        edge[x].push_back(y);
        edge[y].push_back(x);
    }
}

long long dfs(int now,int p,int rt,int mx)
{
    if(mx-a[now]>d || a[now]>mx || (a[now]==mx && p!=-1 && now<rt)) return 0;
    long long res=1;
    for(int i=0;i<edge[now].size();i++)
    {
        int temp=edge[now][i];
        if(temp!=p) res=res*(dfs(temp,now,rt,mx)+1)%MOD;
    }
    return res;
}

void work()
{
    long long ans=0;
    for(int i=1;i<=n;i++) ans=(ans+dfs(i,-1,i,a[i]))%MOD;
    cout<<ans<<endl;
}

int main()
{
    read();
    work();
}