Codeforces 486D Valid Set 树形DP+计数

题意:n个结点的树,每个结点权值为a[i],问有多少个子树其max-min<=d？n,a[i],d<=2e3


直接dp搞不来.. 先不考虑第三个最值条件,可以容易求出以u为根的子树个数 sz[u]=sz[u]*(sz[v]+1)
分类计数:最小值为a[i]的集合个数,以a[i]为根,晒出合法的联通分量 算出sz[u]即可.

注意如果一个联通分量有多个最小值,则会算重复 

当a[u]作为最小值计算完毕后,标记a[u],若之后的i为最小值,a[i]=a[u],i能和u联通,则断开,

因为第一次算a[u]为最小值时 i和u所在的集合已经被计算过一次,所以断开后就不再重复计算.（标记也可以看做编号 规定相同权值点的方向只能从小到大）

#include <bits/stdc++.h>using namespace std;typedef long long ll;const int N=6e3+20;const ll mod=1e9+7;vector<int> e[N];ll d,n,a[N];ll sz[N],ans,mk[N];void dfs(int u,int fa,ll mn){if(a[u]<mn||a[u]>mn+d)return;if(mk[u]&&a[u]==mn)return;if(fa==0)mk[u]=1;sz[u]=1ll;for(int i=0;i<e[u].size();i++){int v=e[u][i];if(v==fa)continue;dfs(v,u,mn);sz[u]=(sz[u]*(sz[v]+1ll)%mod)%mod;}}int main(){while(cin>>d>>n){int u,v;ans=0;for(int i=1;i<=n;i++)scanf("%I64d",&a[i]),e[i].clear();for(int i=1;i<=n-1;i++){scanf("%d%d",&u,&v);e[u].push_back(v);e[v].push_back(u);}memset(mk,0,sizeof(mk));for(int i=1;i<=n;i++)// a[i] being min{memset(sz,0,sizeof(sz));dfs(i,0,a[i]);ans=(ans+sz[i])%mod;//最小值为a[i]的连通分量个数 //cout<<i<<' '<<sz[i]<<endl;}cout<<ans<<endl;}return 0;}