半彩三重奏（可撤销并查集）

半彩三重奏

题目背景

帆帆不满足于只构建一棵深蓝之树，他觉得深蓝之树只有深蓝色太单调了，于是想给这棵树染上颜色。

题目描述

由于现在已经来到了魔幻的龙年，帆帆的深蓝之树已经被染上了颜色，结点 $i$ 的颜色为 $a_i$。

帆帆是一个喜新厌旧的人，在接下来的 $q$ 天中，他每天都会改变他喜欢的颜色，第 $i$ 天他喜欢的两种颜色是 $x_i,y_i$（$x_i\ne y_i$）。

但是为了照顾自己的树，他需要经常在树上移动，并且只会经过自己喜欢的颜色。

具体来说，第 $i$ 天，帆帆会选择一个有序结点对 $(u,v)$，然后沿着 $u\to v$ 的唯一简单路径移动，并且中间经过的结点（包含 $u,v$）颜色必须 ∈ { x i , y i } \in \{x_i,y_i\} ∈{xi​,yi​}（$u$ 可以等于 $v$），之后可以抽取一次宵宫。

每一天你都需要抽取一个满命宵宫，然后告诉帆帆有多少种可能的选择有序结点对的方案。

输入格式

第一行两个正整数 $n,q$。

接下来一行 $n$ 个正整数 $a_1,a_2,\cdots ,a_n$ 表示每个点的颜色。

接下来一行 $n-1$ 个正整数 $p_2,p_3,\cdots ,p_n$，表示对所有 $2\le i\le n$，树上存在一条边 $(p_i,i)$。

接下来 $q$ 行，每行两个正整数 $x_i,y_i$。

输出格式

对于每组询问输出一行一个正整数表示答案。

样例 #1

样例输入 #1

5 3
1 2 1 3 3
1 1 2 2
1 2
1 3
2 3

样例输出 #1

9
6
9

提示

【样例 $1$ 解释】

树的形态如图所示：

对于第一组询问，合法的路径有且仅有 { 1 , 2 , 3 } \{1,2,3\} {1,2,3} 中一点到另一点的路径。

对于第二组询问，合法的路径有且仅有 $1\to 1,1\to 3,3\to 1,3\to 3,4\to 4,5\to 5$。

【子任务约束】

本题采用子任务捆绑测试。

对于 $100\%$ 的数据，有 $1\le n\le 10^6$，$1\le q\le 2\times 10^6$，$1\le a_i,x,y\le n$，$x\ne y$。保证 $p_i<i$。

子任务编号	$n$	$q$	特殊性质	分值	依赖子任务
Subtask #1	$\le 100$	$\le 100$	无	$7$	无
Subtask #2	$\le 2000$	$\le 2000$	无	$18$	$1$
Subtask #3	$\le 10^5$	$\le 2\times 10^5$	A	$5$	无
Subtask #4	$\le 10^5$	$\le 2\times 10^5$	B	$19$	无
Subtask #5	$\le 10^5$	$\le 2\times 10^5$	无	$21$	$1,2,3,4$
Subtask #6	$\le 2\times 10^5$	$\le 5\times 10^5$	无	$10$	$5$
Subtask #7	$\le 10^6$	$\le 2\times 10^6$	无	$20$	$6$

特殊性质 A：$p_i=1$。

特殊性质 B：$p_i=i-1$。

思路

题目大意

• 给定一个 $n$ 个节点的树，每个点 $i$ 上有一个颜色 $a_i$。
• 每次询问，给定两个颜色 $x,y$。你需要找到树上有多少个点对 $(u,v)$，满足 $(u,v)$ 的简单路径上（包括 $u,v$），所有点的颜色要么是 $x$，要么是 $y$。
• $(u,v)$ 和 $(v,u)$ 算两种不同的，$u=v$ 时也算两种不同的。

具体做法

• 首先，这道题比较典型的思路就是：我们把相同颜色的点合并（这个可以用并查集处理）并缩点，然后再按照原图的连线方式将它们连起来。
• 但你如果单纯按照第一点那么做，最慢的地方就在于“拿出树中颜色为 $x$ 和颜色为 $y$ 的节点，进行连边”了。（因为它们会构成森林）
• 所以我们得优化，我们发现：对于树上一条连着两个相同颜色节点的边，它被重复拿出来了好几次的贡献是可以预处理的。具体地，我们先把树上同一颜色并且互相连通的点，缩成一个连通块。这个我们可以用普通并查集来实现。然后我们就可以得到，单独在一个颜色中，可以互达的点对个数。
• 设颜色 $i$ 有 $s_i$ 个连通块 $a_{i,1},a_{i,2},\cdots ,a_{i,s_i}$，每个联通块的大小为 $si{z}_{i,j}$，那么这个颜色中可以互达的点对个数 $inan{s}_i={\sum }_{j=1}^{s_i}si{z}_{i,j}$。这个是算相同颜色的
• 那么不同颜色之间的怎么算呢？
  
  然后，我们处理完一个颜色对之后，我们要把刚才我们所做的连边全部都撤销了。这样，才能保证处理下一个颜色对时不会发生影响。这个可以用栈来写
  
  参考

自己的理解和遇到的问题

• 它那里的答案是 $inan{s}_x+inan{s}_y+stan{s}_{map[(x,y)]}$，我们要特别注意一下，就是它这里的 $inan{s}_x$ 的计算是 ${\sum }_{i=1}^{s}si{z}_i^2$，注意是有平方的，因为你自己模拟一下，同颜色的种类数就是这么多。
• 因为本道题是多组测试数据，最好不要用 $memset$ 很容易超时的。
• 它这边是分为普通并查集和可撤销并查集
  注意，普通并查集的 find这么写：

int find1(int x){
    if(x!=p1[x])p1[x]=find1(p1[x]);
    return p1[x];
}

可撤销的并查集必须得：

int find(int x){
    if(p[x]==x)return x;
	return find(p[x]);
}

这两者的区别就是可撤销并查集不存在路径压缩，所以必须得那么写。（要不然会 $49pts$）

• 如果直接对一个 pair 用 map 进行映射会 TLE，然后 pair 又是不能直接用
  unordered_map 映射的。所以说，我们先把这个 pair (x,y) 给映射成一个整数
  x*10^6+y，再用 unordered_map 进行映射。（这个得特别注意一下）

因为我们规定颜色对 $(x,y)$ 中 $x<y$，所以我们当 $a_x<a_i$ （其中，$a$ 是颜色对），此时的颜色映射为 $tmp=a_x\times 100000+a_i$，反过来就反过来。

• 因为我们要离线处理，因此我们要用 $vp$ 的 $vector$ 去装颜色点对。$vector$ 的索引就是哈希表的值。具体见：

if(!S.count(tmp)){//如果找不到此颜色，则开一个
    S[tmp]=++colp;
}
tp=S[tmp];
ve[tp].push_back({x,i});

• 我们是把普通并查集用于求出各个连通块内部的内容，此时是用来缩点用的。
• 以下是可撤销并查集的模板

void undo(){
    //然后撤销时，直接拿出栈顶的信息，进行复原，再弹栈即可。
    p[stk1.top()]=stk1.top();
    siz[stk2.top().first]=stk2.top().second;
    stk1.pop();
    stk2.pop();
}

for(int i=1;i<=colp;i++){
        for(int j=0;j<ve[i].size();j++){
            int xx=find(ve[i][j].x),yy=find(ve[i][j].y);
            sta[i]+=siz[xx]*siz[yy]*2;//连边
            
            int a=find(ve[i][j].x),b=find(ve[i][j].y);
            if(a!=b){
                if(siz[a]>siz[b])swap(a,b);
                stk1.push(a);//先记录下较小集合的父亲
                stk2.push({b,siz[b]});//再记录下较大集合的父亲和大小
                siz[b]+=siz[a];
                p[a]=b;
                
            }
        }
        for(int j=0;j<ve[i].size();j++)undo();
    }

• 注意答案如果没有颜色对相连（也就是中间有其他颜色），此时答案就不用加上 $stans$ 了。

AC代码

#include<iostream>
#include<algorithm>
#include<cstring>
#include<unordered_map>
#include<vector>
#include<stack>//可撤销并查集

#define int long long
#define x first
#define y second

using namespace std;

typedef pair<int,int>PII;

const int N = 1e6+10;

int a[N],p[N],siz[N],siz1[N];
int p1[N];
int n,m;
int col,tmp;
unordered_map<int,int>S;//颜色对的映射，类似哈希表
struct E{
    int x,y;
}e[N];
vector<E>ve[N];//存储下是 i 颜色对的边有哪些
int insans[N],sta[N];//insans表示颜色中可以互达的点对个数
//sta表示跨颜色之间的贡献
stack<int> stk1;//stk1表示合并的根节点
stack<PII> stk2;//stk2表示已合并到其他根的根大小
int colp,tp;//colp表示图上有多少颜色对在图上出现

int find(int x){
    // if(x!=p[x])p[x]=find(p[x]);
    // return p[x];
    if(p[x]==x)return x;
	return find(p[x]);
}

int find1(int x){
    if(x!=p1[x])p1[x]=find1(p1[x]);
    return p1[x];
}

void undo(){
    //然后撤销时，直接拿出栈顶的信息，进行复原，再弹栈即可。
    p[stk1.top()]=stk1.top();
    siz[stk2.top().x]=stk2.top().y;
    stk1.pop();
    stk2.pop();
}

signed main(){
    scanf("%lld%lld",&n,&m);
    
    for(int i=1;i<=n;i++){
        scanf("%lld",&a[i]);
        p1[i]=i;
    }
    
    for(int i=2;i<=n;i++){
        int x;
        scanf("%lld",&x);
        if(a[x]==a[i]){
            p1[find1(x)]=find1(i);
        }else{
            if(a[x]<a[i]){
                /*如果直接对一个 pair 用 map 进行映射会 TLE，然后 pair 又是不能直接用
                unordered_map 映射的。所以说，我们先把这个 pair (x,y) 给映射成一个整数 
                x*10^6+y，再用 unordered_map 进行映射。
                */
                tmp=a[x]*1e6+a[i];
            }else{
                tmp=a[i]*1e6+a[x];
            }
            
            if(!S.count(tmp)){//如果找不到此颜色，则开一个
                S[tmp]=++colp;
            }
            tp=S[tmp];
            ve[tp].push_back({x,i});
        }
    }
    
    
    for(int i=1;i<=n;i++){
        p1[i]=find1(i);
        siz1[p1[i]]++;
    }
    
    
    for(int i=1;i<=n;i++){
        if(p1[i]==i){
            insans[a[i]]+=siz1[i]*siz1[i];
        }
    }
    
    for(int i=1;i<=n;i++)p[i]=p1[i];
    
    for(int i=1;i<=n;i++)siz[i]=siz1[p1[i]];
    //前面的过程为预处理各个颜色缩点连线处理
    
    for(int i=1;i<=colp;i++){
        for(int j=0;j<ve[i].size();j++){
            int xx=find(ve[i][j].x),yy=find(ve[i][j].y);
            sta[i]+=siz[xx]*siz[yy]*2;//连边
            
            int a=find(ve[i][j].x),b=find(ve[i][j].y);
            if(a!=b){
                if(siz[a]>siz[b])swap(a,b);
                stk1.push(a);//先记录下较小集合的父亲
                stk2.push({b,siz[b]});//再记录下较大集合的父亲和大小
                siz[b]+=siz[a];
                p[a]=b;
                
            }
        }
        for(int j=0;j<ve[i].size();j++)undo();
    }
    
    
    while(m--){
        int x,y;
        scanf("%lld%lld",&x,&y);
        int ans=insans[x]+insans[y];
        if(x>y)swap(x,y);
        tmp=x*1e6+y;
        if(!S.count(tmp)){
            printf("%lld\n",ans);
        }else{
            //没有边连接
            tp=S[tmp];
            ans+=sta[tp];
            printf("%lld\n",ans);
        }
    }
    return 0;
}