Nowcoder 5278G.血压游戏 第18届上海大学网络友谊赛（虚树+dp）

题目描述

Compute 有一棵 n 个点，编号分别为 1∼n 的树，其中 s 号点为根。
Compute 在树上养了很多松鼠，在第 i 个点上住了 ai个松鼠。

因为某些缘故，它们开始同时向根节点移动，但它们相当不安分，如果在同一个节点上，它们就会打起来，简单地来说以下事件会依序发生：

·如果一个节点上有 2 只或 2 只以上的松鼠，他们会打架，然后这个节点上松鼠的数量会减少 1；

·根节点的所有松鼠移动到地面，位于地面上的松鼠不会再打架；

·所有松鼠同时朝它们的父节点移动。

所有事件各自都在一瞬间完成，直至树上没有松鼠。

现在 Compute 想知道最终有多少只松鼠到达了地面。

输入描述:

第一行包含两个整数 n, s (1≤n≤2e5,1≤s≤n)，中间以空格分隔，分别表示点的数量和根的编号。

第二行包含 n 个整数 a1, a2, an(0≤ai≤1e9)，中间以空格分隔，分别表示每个点上一开始的松鼠数量。

接下来 n-1 行，每行包含两个整数 u, v (u!=v, 1≤u,v≤n)，中间以空格分隔，表示 u 和 v 之间有一条边。

输入保证是一棵树。

输出描述:

在一行输出一个整数，表示最终到达地面的松鼠数量。

示例1

输入

3 1
2 4 6
1 2
1 3

输出

8

示例2

输入

3 1
0 1 1
1 2
1 3

输出

1

算法标签：虚树、dp

首先读完题发现，可以将这棵树逐层分离，深度不同的结点不会互相影响，因此只需要枚举层数，每次取出深度相同的点做一次dp就可以了。然后先总结自己的一个错误吧。训练的时候当然先考虑暴力的复杂度，就是O(max{deep}*n)。然后很理想地构造了一个满二叉树，发现复杂度就是O(nlogn)了。。然后居然就觉得暴力是可做的。

然而其实只要一个数据点给一条链就退化成O(n^2)了。其实做的时候也知道，只是抱侥幸心理。但是树形题的数据怎么可能没有链呢。。 所以以后不可以这样了.jpg

训练完看题解才知道要用虚树做，于是正好回过头借这题补一下虚树的知识。当然建虚树的算法基于dfs序和lca（最近共同祖先），这里就先不展开，重点放在虚树上面。

虚树的基本思想就是从原来完整的树上取出需要的结点，同时取出任意两个结点的lca和根节点，在保证原先的祖先——后代关系的前提下，将这些点再连成一个新树。比如下图，我想要取出左图红色的三个点，那么就需要取出它们和它们两两的lca，建成右图的树状关系。

这样的好处很明显，如果打绿叉的点或链对答案无关紧要（或者可以通过O(1)快速计算），我们直接在虚树上进行计算的速度会远快于原来完整的树。

具体会快多少其实不太好算，但是我们知道k个结点两两之间的lca最多有k-1个（极端情况如图）,那么虚树中结点最多是2k个。所以即使考虑最坏的情况，对整个树做dp的复杂度也是O(n)级别的。

当然这里还没有考虑建树的复杂度。那么怎么建虚树呢？朴素思想当然是枚举任意两个点，将这些点和它们的lca都放入集合里，再根据祖先——后代关系连边。但是想想就知道太慢了。。所以就需要dfs序，并引入一个单调堆栈sta。这里的单调性指的是确保堆栈中的元素都在原树的一条链上，并且自底向上保持祖孙关系。比如说堆栈的情况如左图所示，那么可以得出虚树结点间的关系如右图。

第二个问题是怎样插入结点呢。如果当前将要插入的结点是5号结点，且5号结点是2号结点的后代，则只要直接将5号加入堆栈即可。

但如果5号结点不是2号结点的子孙呢？这时候事情就会变得比较麻烦。比如当5号结点是3号结点的右儿子时，直接将5号加入堆栈显然不能保持堆栈的单调性。正确的做法是先依次将2号、1号、4号结点弹出堆栈（在出栈的同时连边），再将5号加入堆栈。

当然，这样讲只是便于一个直观的理解，因为还有许多没有考虑周全的情况。比如怎样确保链的有序性？如果祖先结点还未在虚树中，应该如何同时添加两个结点？解决这些问题的答案是dfs序。严谨的构建方法如下：

我们需要将准备插入的点根据dfs序从小到大排序（令结点n的dfs序为dfn[n]），如果当前要加入的点为x，当前栈顶为y，计算出x和y的lca(x,y)，此时分两种情况：
1. lca=y， 意味着x是y的子孙，直接将x入栈即可
2. lca!=y，说明x和y分别在lca的两棵子树中，此时y所在的子树已经构建完毕（由dfs序的性质可知，假设y所在的子树里还有结点k未加入，则dfn[k] < dfn[x]，应该先访问k），此时需要同时进行出栈和连边的操作：

while(TOP>=2 && dfn[sta[TOP-1]]>=dfn[lca]){
    E[sta[TOP-1]].push_back(sta[TOP]);
    TOP--;
}

一开始写这段代码的时候不太理解为什么要将lca与堆栈的第二个元素sta[TOP-1]而不是栈顶元素sta[TOP]进行比较，可以考虑下图的情形：

此时最关键的问题是lca并不在栈中，且dfs序的关系有sta[TOP-1] < lca < sta[TOP]。因此如果将lca与栈顶元素比较，while条件成立，会连一条蓝点sta[TOP-1]->绿点sta[TOP]的边，但明显我们需要的是lca->sta[TOP]这条边。而如果和第二个元素比较，进行到这一步时while循坏就停止了，我们只需要再额外加一个判断，连上lca->sta[TOP]，再将栈顶元素替换成lca即可。

if (sta[TOP]!=lca) E[lca].push_back(sta[TOP]),sta[TOP]=lca;

此时y所在的子树所有边都以连好并全部出栈，lca也已入栈，最后只需要将x入栈，该结点的插入操作就完成了。

sta[++TOP]=x;

最后当所有结点都插入完成后，清空堆栈并连边，建出最后一条链。

简而言之，虚树的构建是以链为单位，从左至右进行，并由dfs序保证其有序性。

dp部分比较简单，也不是这篇的重点，就直接放在代码里了。感觉这题还挺适合做模板的（dfs序、lca、建虚树），所以也借鉴了一些别人的代码，专门把各个部分优化了一下。（主要是lca部分，对lg数组的预处理）也算是存个档吧。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=2e5+5;

int pre[N][25],dfn[N],deep[N],sta[N],TOP,lg[N];

int A[N],n,s,NODE,x,y;

ll ans=0;

vector<int>G[N];
vector<int>E[N];
vector<int>D[N];

bool cmp(int x,int y){ return dfn[x]<dfn[y]; }

void DFS(int x,int fa)
{
    int to;
    dfn[x]=++NODE;
    deep[x]=deep[fa]+1;
    D[deep[x]].push_back(x);
    pre[x][0]=fa;
    for (int i=1;(1<<i)<=deep[x];i++) pre[x][i]=pre[pre[x][i-1]][i-1];
    for (int i=0;i<G[x].size();i++)
    {
        to=G[x][i];
        if (to!=fa)
            DFS(to,x);
    }
}

int LCA(int x,int y)
{
    if (deep[x]<deep[y]) swap(x,y);
    while(deep[x]>deep[y]) x=pre[x][lg[deep[x]-deep[y]]-1];
    if (x==y) return x;
    for (int i=lg[deep[x]]-1;i>=0;i--)
        if (pre[x][i]!=pre[y][i])
            x=pre[x][i],y=pre[y][i];
    return pre[x][0];
}

void Insert(int x)
{
    if (TOP==1){
        if (sta[TOP]!=x) sta[++TOP]=x;
        return;
    }
    int lca=LCA(x,sta[TOP]);
    while(TOP>=2 && dfn[sta[TOP-1]]>=dfn[lca]){
        E[sta[TOP-1]].push_back(sta[TOP]);
        TOP--;
    }
    if (sta[TOP]!=lca) E[lca].push_back(sta[TOP]),sta[TOP]=lca;
    sta[++TOP]=x;
}

ll solve(int x)
{
    if (E[x].empty())
        return 1LL*A[x];
    ll ret=0;
    for (int i=0;i<E[x].size();i++)
    {
        int to=E[x][i];
        ll sol=solve(to);
        if (sol!=0)
            ret+=max(1LL,sol-(deep[to]-deep[x]));
    }
    E[x].clear();
    return ret;
}

int main()
{
    cin>>n>>s;
    for (int i=1;i<=n;i++) lg[i]=lg[i-1]+(1<<(lg[i-1])==i);
    for (int i=1;i<=n;i++) scanf("%d",&A[i]);
    for (int i=1;i<=n-1;i++)
    {
        scanf("%d%d",&x,&y);
        G[x].push_back(y);
        G[y].push_back(x);
    }
    DFS(s,0);
    for (int i=1;!D[i].empty();i++)
    {
        sort(D[i].begin(),D[i].end(),cmp);
        TOP=0;
        sta[++TOP]=s;
        for (int j=0;j<D[i].size();j++)
        {
            Insert(D[i][j]);
        }
        while(TOP>1)
        {
            E[sta[TOP-1]].push_back(sta[TOP]);
            TOP--;
        }
        ll sol=solve(s);
        if (sol) ans+=max(1LL,sol-1);
    }
    printf("%lld",ans);
    return 0;
}