NC19314 颓红警(dfs+思维)

这个depth的式子就很有那味，明显就是把他拆分掉，维护0次项，一次项和二次项

因为只能从头，所以递归维护即可，注意到如果小于0了，就把他的贡献剪掉。

我第一次写的时候用了dfs序+线段树维护，复杂度感觉还行，但是内存和常数太大，过不了，只过了70%，这题本来也不用这么麻烦做

数据量给到1e6，用大常数log还是比较虚的

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<ll,ll> pll;
const int N=2e6+10;
const int mod=1e9+7;
ll n,p,lim;
int pa[N];
ll ans;
ll s1[N],s2[N],s[N],cnt[N],a[N],tot[N];
ll h[N],ne[N],e[N],idx;
void add(int a,int b){
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
void dfs(int u,int fa,ll d){
    if(d>lim){
        ll t=d-lim;
        s[u]-=cnt[t];
        s1[u]-=cnt[t]*t;
        s2[u]-=cnt[t]*t*t;
    }
    ll tmp=s[u]*(p-d*d)-s2[u]+2*d*s1[u];
    if(a[u]>=tmp)
        tot[u]=(a[u]-tmp)/p+1,ans+=tot[u];
    cnt[d]=tot[u];
    for(int i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa)
            continue;
        s[j]=s[u]+tot[u];
        s1[j]=s1[u]+tot[u]*d;
        s2[j]=s2[u]+tot[u]*d*d;
        dfs(j,u,d+1);
    }
}
int main(){
    //ios::sync_with_stdio(false);
    memset(h,-1,sizeof h);
    cin>>n>>p;
    int i;
    lim=sqrt(p-1)+1;
    for(i=1;i<=n;i++){
        scanf("%d",&a[i]);
    }
    for(i=1;i<n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
        pa[y]=x;
    }
    int rt;
    for(i=1;i<=n;i++){
        if(pa[i]==0){
            rt=i;
            break;
        }
    }
    dfs(rt,-1,1);
    cout<<ans<<endl;
    return 0;
}