Description：

给出一棵 $n$ 个节点的树，每个点有颜色 $c_i$ ，求树上有多少个连通块包含不超过两种不同颜色。

$1\leqslant n\leqslant 5\times 10^5$

Solution：

显然是联通块 $DP$ ，首先假设根节点一定会被包含，那么根节点的颜色是一定包含的，那么其实我们只要在状态里记录另外一个颜色就行了，也就是设 $f[i][j]$ 表示以 $i$ 为根的子树，另外一种颜色是 $j$ 的联通块数，这样不难写出一个 $O(n^2)$ 的 $DP$ ，观察一下这个 $DP$ 需要这样几个操作：全局赋 $1$ ，和子树合并对应位置相乘，单点求值，单点赋值，全局求和，显然这些都可以用线段树做，也就是整体 $DP$ 。

Code：

#include<algorithm>

#include<iostream>

#include<cstdlib>

#include<cstdio>

#include<cmath>

#include<cctype>

#include<cstring>

using namespace std;

int n;

#define MAXN 300010

#define N 300000

#define MOD 1000000007

inline int rd()

{

	register int res = 0,f = 1;register char c = getchar();

	while(!isdigit(c)){if(c == '-')f = -1;c = getchar();}

	while(isdigit(c))res = (res << 1) + (res << 3) + c - '0',c = getchar();

	return res * f;

}

struct edge{int to,nxt;}e[MAXN << 1];

int edgenum = 0,lin[MAXN] = {0};

void add(int a,int b)

{

	e[++edgenum] = (edge){b,lin[a]};lin[a] = edgenum;

	e[++edgenum] = (edge){a,lin[b]};lin[b] = edgenum;

}

int col[MAXN];

struct node

{

	int lc,rc;

	int sum,add,mul;

	node(){sum = add = 0;mul = 1;}

}t[MAXN * 20];

int ptr = 0;

int newnode(){return ++ptr;}

int root[MAXN];

#define mid ((l + r) >> 1)

void pushdown(int rt,int l,int r)

{

	int ll = l,lr = mid,rl = mid + 1,rr = r;

	if(t[rt].lc == 0)t[rt].lc = newnode();

	if(t[rt].rc == 0)t[rt].rc = newnode();

	t[t[rt].lc].sum = (1ll * t[t[rt].lc].sum * t[rt].mul % MOD + 1ll * (lr - ll + 1) * t[rt].add % MOD) % MOD;

	t[t[rt].rc].sum = (1ll * t[t[rt].rc].sum * t[rt].mul % MOD + 1ll * (rr - rl + 1) * t[rt].add % MOD) % MOD;

	t[t[rt].lc].mul = 1ll * t[t[rt].lc].mul * t[rt].mul % MOD;

	t[t[rt].rc].mul = 1ll * t[t[rt].rc].mul * t[rt].mul % MOD;

	t[t[rt].lc].add = (1ll * t[t[rt].lc].add * t[rt].mul % MOD + t[rt].add) % MOD;

	t[t[rt].rc].add = (1ll * t[t[rt].rc].add * t[rt].mul % MOD + t[rt].add) % MOD;

	t[rt].mul = 1;t[rt].add = 0;

	return;

}

int merge(int x,int y,int l,int r)

{

	if(x == 0 || y == 0)return x + y;

	if(l == r)

	{

		t[x].sum = 1ll * t[x].sum * t[y].sum % MOD;

		return x;

	}

	if(t[x].lc == 0 && t[x].rc == 0)swap(x,y);

	if(t[y].lc == 0 && t[y].rc == 0)

	{

		t[x].sum = 1ll * t[x].sum * t[y].add % MOD;

		t[x].mul = 1ll * t[x].mul * t[y].add % MOD;

		t[x].add = 1ll * t[x].add * t[y].add % MOD;

		return x;

	}

	pushdown(x,l,r);pushdown(y,l,r);

	t[x].lc = merge(t[x].lc,t[y].lc,l,mid);

	t[x].rc = merge(t[x].rc,t[y].rc,mid + 1,r);

	t[x].sum = (t[t[x].lc].sum + t[t[x].rc].sum) % MOD;

	return x;

}

int query(int rt,int p,int l,int r)

{

	if(l == r)return t[rt].sum;

	if(t[rt].lc == 0 && t[rt].rc == 0)return t[rt].add;

	pushdown(rt,l,r);

	if(p <= mid)return query(t[rt].lc,p,l,mid);

	else return query(t[rt].rc,p,mid + 1,r);

}

void set(int rt,int p,int k,int l,int r)

{

	if(l == r){t[rt].sum = k;return;}

	pushdown(rt,l,r);

	if(p <= mid)set(t[rt].lc,p,k,l,mid);

	else set(t[rt].rc,p,k,mid + 1,r);

	t[rt].sum = (t[t[rt].lc].sum + t[t[rt].rc].sum) % MOD;

	return;

}

int ans = 0;

void dp(int k,int fa)

{

	root[k] = newnode();t[root[k]].sum += N + 1;t[root[k]].add += 1;

	for(int i = lin[k];i != 0;i = e[i].nxt)

	{

		if(e[i].to == fa)continue;

		dp(e[i].to,k);

		if(col[k] == col[e[i].to])

		{

			t[root[e[i].to]].sum += N + 1;t[root[e[i].to]].add += 1;

			root[k] = merge(root[k],root[e[i].to],0,N);

		}

		else

		{

			set(root[k],col[e[i].to],1ll * query(root[k],col[e[i].to],0,N) * (query(root[e[i].to],col[k],0,N) + 1) % MOD,0,N);

		}

	}

	ans = (ans + t[root[k]].sum) % MOD;

	ans = (ans - 1ll * query(root[k],0,0,N) * N % MOD + MOD) % MOD;

	return;

}

int main()

{

	scanf("%d",&n);

	for(int i = 1;i <= n;++i)col[i] = rd();

	for(int i = 2;i <= n;++i)add(i,rd());

	dp(1,0);

	cout << ans << endl;

	fclose(stdin);

	fclose(stdout);

	return 0;

}