CodeForces 486D vailid set 树形DP

题意：

给一颗树，在树中寻找所有符合点权最大值和最小值相减小于d的联通子集。 

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll mod=1e9+7;
const int  maxn=4e3+10;
struct node{
    int to,next;
}e[maxn];
int head[maxn],a[maxn];
int d,n;
int flag=0;
void add(int u,int v){
    e[flag].to=v;
    e[flag].next=head[u];
    head[u]=flag++;
}
//void add(int u,int v){
    e[++tot].to=u;
    e[tot].next=head[v];
    head[v]=tot;
//    e[++tot].to=v;
//    e[tot].next=head[u];
//    head[u]=tot;
//}
ll dfs(int x,int pre,int rt){
    ll ans=1;
    for(int i=head[x];i!=-1;i=e[i].next){
        //前一个点不等于下一个点//下一个点点权大于根或者(等于根且编号小于根）//下一个点的点权和根的差值小于d
        if(pre!=e[i].to&& (a[e[i].to]>a[rt]||(a[e[i].to]==a[rt]&&rt>e[i].to)) && a[e[i].to]-a[rt]<=d ){
            ans=(ans*(dfs(e[i].to,x,rt)+1)%mod)%mod;
        }
    }
    return ans;
}
int main(){
    scanf("%d%d",&d,&n);
    memset(head,-1,sizeof(head));
    for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
    }
    for(int i=1;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v);add(v,u);
    }
//    cout<<"tet"<<endl;
    ll ans=0;
    for(int i=1;i<=n;i++){
//        cout<<dfs(i,0,i)<<endl;
        ans+=dfs(i,0,i);
    }
    printf("%lld",ans%mod);
}