NC19314 颓红警(DFS+树形DP)

题目链接

题意：
$有一棵树代表敌人的分级网络，n个结点$
$从上到下依次等级地位下降，每个结点代表一个部队$
$每个部队有攻击力a_i，每次可以使一个结点的攻击力减少p$
$同时这个结点的子节点会减少p-dis(i,j{)}^2$
$如果攻击力小于0，那么将被歼灭$
$只能对未被歼灭的部队发动进攻$
$问最多多少次能够歼灭这个军队$
题解：
$n<=1e6,a_i<=1e9$
$由于是个树，而且儿子和父亲有关系$
$最先考虑的就是DFS遍历然后进行状态转移$
$如果攻击父结点，就会对子结点造成伤害$
$这个伤害值为p-dis(i,j{)}^2$
$对这个进行化简，可以得到p-(d_j-d_i{)}^2$
$(d表示结点的深度)$
$所以现在需要看的就是每个父亲对该结点的影响$
$通过这个式子我们可以发现超过父结点只能对一定距离的子节点造成伤害$
$这个距离就是\sqrt{p}，对于每个儿子将多余的影响减去$
$每个儿子最多收到这么多深度的影响$
$然后把上述式子化简$
$p-(d_j^2-2\ast d_i\ast d_j+d_i^2)$
$这样就发现，其实对于每个子结点$
$只和上述结点的攻击次数，深度和深度的平方有关系$
$所以直接用树形dp维护这些值$
$但是由于刚才算的最远距离，所以每次超出这个距离的时候需要减去$
$每次贪心对尽量靠上的结点进行攻击$
$最后计算攻击次数$
AC代码

/*
    Author:zzugzx
    Lang:C++
    Blog:blog.csdn.net/qq_43756519
*/
#include<bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(),(x).end()
#define endl '\n'
#define SZ(x) (int)x.size()
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const int mod=1e9+7;
//const int mod=998244353;
const double eps = 1e-10;
const double pi=acos(-1.0);
const int maxn=1e6+10;
const ll inf=0x3f3f3f3f;
const int dir[][2]={{0,1},{1,0},{0,-1},{-1,0},{1,1},{1,-1},{-1,1},{-1,-1}};

ll ans,n,p,lim;
vector<int> g[maxn];
ll s1[maxn],s2[maxn],s[maxn],cnt[maxn],a[maxn],tot[maxn];
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(auto v:g[u]){
        if(v==fa)continue;
        s[v]=s[u]+tot[u];
        s1[v]=s1[u]+tot[u]*d;
        s2[v]=s2[u]+tot[u]*d*d;
        dfs(v,u,d+1);
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    //freopen("in.txt","r",stdin);
    //freopen("out.txt","w",stdout);
    cin>>n>>p;lim=sqrt(p-1)+1;
    for(int i=1;i<=n;i++)cin>>a[i];
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs(1,0,1);
    cout<<ans;
    return 0;
}