2020暑期牛客多校训练营第七场（F）Tokens on the Tree（树链剖分，思维，树形dp）

Tokens on the Tree

原题请看这里

题目描述：

$Chiaki$有一棵 $n$ 个顶点的树，树的每个顶点可能被标记为白色或者黑色，有 $w$ 个白点和 $b$ 个黑点。对于颜色相同的每对顶点，它们之间必须存在一条路径，路径上的每个顶点包含颜色标记，且颜色相同。
$Chiaki$希望执行以下操作：

• 选择一个带有标记的顶点。
• 选择一个路径 $p_1,p_2,...,p_k$，其中 $p_1=u$，所有顶点 $p_1,p_2,...,p_{k-1}$都包含相同颜色的标记，而pk中没有标记。
• 将 $p_1$中的标记移至 $p_k$。现在 $p_1$中没有颜色标记，而 $p_k$有颜色标记。

对于标记 $S$和 $T$的两种初始配置，如果 $Chiaki$可以多次执行上述操作以使 $S$变为 $T$，则认为 $S$和 $T$是等价的。
令 $f(w,b)$为等价类的数量（等价类是其中元素相互等价的最大配置集；所有等价类对所有配置的集合进行划分）， $Chiaki$希望知道
$(\mathop{\sum}\limits_{w=1}^{n-1}\mathop{\sum}\limits_{b=1}^{n-w}w*b*f(w,b))$ $mod$ $(10^9+7)$的值。

输入描述：

有多个测试用例。 输入的第一行包含一个整数 $T$，指示测试用例的数量。
对于每个测试用例，第一行包含一个整数 $n(2 \le n \le 2 \times 10 ^ 5)$-树中的顶点数。 第二行包含 $n-1$个整数 $p_2，p_3，\dots，p_n(1 \le p_i < i)$，其中 $p_i$表示顶点 $i$和顶点 $p_i$之间存在边，可以保证所有测试用例的 $n$之和 不会超过 $2 \times 10 ^ 5$

输出描述：

对于每个测试用例，输出一个整数，该整数表示 $(\mathop{\sum}\limits_{w=1}^{n-1}\mathop{\sum}\limits_{b=1}^{n-w}w*b*f(w,b))$ $mod$ $(10^9+7)$的值。

样例输入：

2
5
1 2 3 3
10
1 2 3 4 3 6 3 8 2

样例输出：

71
989

说明：

在第一组数据中,f(w,b)的值对于每个w和b: 
f(1,1)=1,
f(1,2)=2, 
f(1,3)=3, 
f(1,4)=3, 
f(2,1)=2,
f(2,2)=2, 
f(2,3)=1, 
f(3,1)=3, 
f(3,2)=1, 
f(4,1)=3.

思路：

首先我们分析一下题目所给的操作，我们发现其实就是可以在树上整块整块的移动 （蠕动） 的两个连通块，相同的颜色在同一块区域里移动只能算一种情况。
接下来我们看一个简单的例子：
在这个例子中，我们发现红色部分占据了图中很大一部分，包括了中间的紫色关键点，无论另一种颜色在蓝色部分还是黄色部分都无法移动。这时我们就相对好办了，只要找到中间紫色的关键点，再跑树形dp就可以了。
但是还有一个麻烦的例子：

在这个例子中，我们发现蓝色连通块可以为所欲为，哪里都可以到，但是如果红色的又堵住了蓝色的路，就又没有地方可以去了，这种情况就比较难实现。
这时，我们就可以用到树链剖分了。
我们将蓝色的当做障碍，如果红色的要到另一边去，那么蓝色连通块需要至少让出一个节点使红色的能够到达另一边，比如上图，蓝色的显得很安稳，让路了。

但是在上图中，蓝色连通块就挡住了红色的去路。
我们使图中较大的一个连通块为红色，另一个蓝色，那么红色肯定要从左端到右端，那么问题来了，如何定义这个左右呢？我们就需要求出图中的最长链，令一端为左，一端为右即可。那么如何求树的最长链？只要求出以重心为根的最长链和次长链，再连起来就可以了。
具体细节请看代码。

$AC$ $Code$:

#include<bits/stdc++.h>
#define ll long long
#define P pair<int,int>
using namespace std;
const int mod=1e9+7;
const int MAXN=2e5+5;
pair<int,int> a;
vector<int> vec[MAXN],tor;
int nd[MAXN],dep[MAXN],vl[MAXN],vr[MAXN],vc[MAXN];
//nd[]表示当前节点为根节点的子树大小
//dep[]表示当前节点深度
//vl[i] 表示以长重链上第i个节点左侧节点（第i-1个节点）为根节点的子树大小
//vr[i] 表示以长重链上第i个节点右侧节点（第i+1个节点）为根节点的子树大小
//vc[i] 表示以与长重链上第i个节点相邻，但不在长重链上的节点为根节点的最大子树大小
//pair型的a：最大子节点的子树的大小，重心id
//tor存长重链（最长重链+次长重链）
ll ans;
int n,t,q;
int getsize(int pos,int p){//求p为v的父节点时，以v为根节点的子树大小
    if(dep[pos]>dep[p]) return nd[pos];
    return n-nd[p];
}
ll getsum(ll x,ll y){//等差数列求和，求x+(x+1)+(x+2)+...+(y-1)+y
	return (x+y)*(y-x+1)/2%mod;
}
void dfs1(int pos,int p,int deep){//寻找树的重心
    nd[pos]=1;
    dep[pos]=deep;
    int maxn=0;
    for(int i=0;i<vec[pos].size();++i){
        int v=vec[pos][i];
        if(v^p){
            dfs1(v,pos,deep+1);
            nd[pos]+=nd[v];
            maxn=max(maxn,nd[v]);
        }
    }
    maxn=max(maxn,n-nd[pos]);
    a=min(a,P(maxn,pos));
}
void dfs2(int pos,int p){//从重心出发的重链，存入tor中
    tor.push_back(pos);
    int maxn=0;
    for(int i=0;i<vec[pos].size();++i){
        int v=vec[pos][i];
        if(v^p) maxn=max(maxn,getsize(v,pos));
    }
    for(int i=0;i<vec[pos].size();++i){
        int v=vec[pos][i];
        if(v^p&&getsize(v,pos)==maxn){
        	dfs2(v,pos);
        	break;
        }
    }
}
void dfs3(int pos,int p,int x){
    int y=getsize(p,pos),z=n-y;
//y表示除v子树外，其他点的数量
//z表示v子树大小
    ans=(ans+getsum(x,y)*getsum(1,z)%mod*2ll%mod)%mod;
//此时x<=big<=y,1<=samll<=z，有f(big,small)=f(small,big)=1两种情况
//贡献=(x+(x+1)+...+(y-1)+y)*(1+2+...+z)*2
    for(int i=0;i<vec[pos].size();++i){//枚举small颜色出现在的子树
        int v=vec[pos][i];
        if(v^p) dfs3(v,pos,getsize(p,pos)+1);
//将big继续增加,递归求解。此时big颜色至少覆盖v节点
    }
}
int main(){
    scanf("%d",&t);
    while(t--){
        scanf("%d",&n);
        for(int i=0;i<n;++i)vec[i].clear();
        for(int i=1;i<n;++i){
            scanf("%d",&q);q--;
            vec[i].push_back(q);
            vec[q].push_back(i);
        }
		a=P(n,-1);
        dfs1(0,-1,0);//求树的重心
        int root=a.second;
        tor.clear();
        dfs2(root,-1);//求从重心出发的重链
        reverse(tor.begin(),tor.end());
        tor.pop_back();//翻转，弹出重心。得到除去重心的第一重链
        dfs2(root,tor[tor.size()-1]);//将从重心出发，第二重的链也存入tor中
//两次dfs2后，tor中存储第一重链+第二重链组成的长重链
        multiset<int> st;
        multiset<int>::iterator it;
        for(int i=0;i<tor.size();++i){//从第一重链的尾部出发，到第二重链的尾部停止
            int pos=tor[i];
            vl[i]=vc[i]=vr[i]=0;
            for(int j=0;j<vec[pos].size();++j){
                int v=vec[pos][j];
                if(i&&tor[i-1]==v) vr[i]=n-getsize(v,pos);//重链上i相邻节点的子树的节点数
                else if(i+1<tor.size()&&tor[i+1]==v) vl[i]=n-getsize(v,pos);//重链上i另一个相邻的节点的子树的节点数
                else vc[i]=max(vc[i],getsize(v,pos));//重链外i相邻节点的最大的子树节点数
            }
            st.insert(vc[i]);//用于存储 vc[]
        }
        int l=0,r=tor.size()-1,big=1;
        ans=0;
        for(;;big++){//枚举w和b中的较大值big,设另一个为samll
//设长重链的左右两端分别放置big对应的颜色,考虑small颜色的所有可能情况
            while(l<r&&vl[l]<big){//删除长重链两端的重子树小于bi 的节点及其相关 vc[] 数据
                l++;
                if(l==r) break;
                it=st.lower_bound(vc[l]);
                st.erase(it);
            }while(l<r&&vr[r]<big){//同上
                r--;
                if(l==r) break;
                it=st.lower_bound(vc[r]);
                st.erase(it);
            }
            if(l>=r) break;//长重链为空只有一个重心时，结束
            int maxn=0;
            if(!st.empty()){//找到剩余的集合中的最大值，用于放置small作为中转
                it=st.end();
                it--;
                maxn=*it;
            }
            ans=(ans+(ll)big*(ll)big%mod*(maxn>=big?1ll:2ll)%mod)%mod;
//big=samll时的情况。若mx>=small=big，则big颜色和samll颜色可以调换位置 f(small,big)=1，否则 f(small,big)=2
            int small=min(maxn,big-1);
//small<big且small<=mx 时的情况，此时两种颜色可以任意交换位置，f(small,big)=1
            ans=(ans+(ll)big*(ll)small%mod*(ll)(small+1)%mod)%mod;
//有f(small,now)和f(now,small)两种情况，总贡献=now*(1+2+...+small)*2*f(big,small)
            small=big-1;
            if(maxn<small)//mx<small<big时的情况，此时两种颜色不能交换位置，small颜色在big颜色的两边之一，f(small,big)=2
//同样有f(big,small)和f(small,big)两种情况，=now*((mx+1)+(mx+2)+...+to)*2*f(small,big)
                ans=(ans+(ll)big*(ll)(small+maxn+1)%mod*(ll)(small-maxn)%mod*2ll%mod)%mod;
        }
        for(int i=0;i<vec[root].size();++i){//遍历重心的相邻节点，此时重心必定被big颜色覆盖，big和small颜色不再能交换位置
            int v=vec[root][i];
            dfs3(v,root,big);//考虑big继续增大的情况
        }
        printf("%lld\n",ans);
    }
}