Codeforces 461B - Appleman and Tree 树状DP

一棵树上有K个黑色节点，剩余节点都为白色，将其划分成K个子树，使得每棵树上都只有1个黑色节点，共有多少种划分方案。

个人感觉这题比较难。假设dp(i,0..1)代表的是以i为根节点的子树种有0..1个黑色节点的划分方案数。

当节点i为白色时，对于它的每个孩子的节点处理：

求dp(i, 0)时有：

1，将该节点与孩子节点相连，但要保证孩子节点所在的子树种没有黑色节点；

2，将该节点不与该孩子节点相连，则该孩子节点要保证所在子树种有黑色节点；

即dp(i, 0) = π(dp(j,0 ) + dp(j, 1)) ，其中j为i的孩子节点

求dp(i,1)时有：

将该节点与其中每个孩子节点中的一个相连，并且保证该孩子节点所在子树中有1个黑色节点（所以共有K种情况，K为该节点的孩子数），并且对于剩下的节点可以选择连也可以选择不连，如果连接，则保证该子节点所在子树中没有黑色，如果不连，则要保证有黑色，所以对于剩下的每个

子节点的处理方案书有dp(j,0) + dp(j,1)个，然后将每个孩子处理的方案书相乘即可，最后将所有的方案相加即可。

当节点i为黑色的时候，求dp(i, 0) 肯定是0；

求dp(i, 1)时对于i的每个子节点也是有两种选择，连或者不连，如果连接，则保证该子节点所在子树中没有黑色，如果不连，则要保证有黑色，即对于每个子节点的处理数共有

dp(j, 0) + dp(j, 1)个，然后将每个孩子处理的方案数相乘。

最终dp(0,1)即为答案，这里假设0节点为根节点。

过程中可以加个小小的优化，当一个子节点所在的整棵子树中若没有黑色节点，那么该节点肯定与其父节点相连，所以计算时可以不考虑该节点。

#include <stdlib.h>
#include <stdio.h>
#include <algorithm>
#include <vector>

using namespace std;
//int values[500001];
//long long sums[500001];
#define MODVALUE 1000000007
#define MOD(x) if((x) > MODVALUE) x %= MODVALUE;

struct Edge
{
	int to;
	int i;
	int totalcolor;
	Edge()
	{
		totalcolor = 0;
	}
};

int compp(const void* a1, const void* a2)
{
	return *((int*)a2) - *((int*)a1);
}

vector<Edge> G[100001];
int Color[100001];
long long res[100001][2];
//int TMP[100001];
bool Visited[100001];

void AddEdge(int from, int to)
{
	Edge edge;
	edge.to = to;

	edge.i = G[to].size();
	G[from].push_back(edge);
	edge.to = from;

	edge.i = G[from].size() - 1;
	G[to].push_back(edge);

}

int CountColor(int node)
{
	Visited[node] = true;
	int count = 0;
	if (Color[node])
	{
		count = 1;
	}
	for (int i = 0; i < G[node].size();i++)
	{
		Edge& edge = G[node][i];
		if (!Visited[edge.to])
		{
			edge.totalcolor = CountColor(edge.to);
			count += edge.totalcolor;
		}

	}
	return count;
}

void GetAns(int node)
{
	Visited[node] = true;
	long long ans = 1;
	int countofcolor = 0;
	vector<int> TMP;
	for (int i = 0; i < G[node].size(); i++)
	{
		Edge& edge = G[node][i];
		if (Visited[edge.to])
		{
			continue;
		}
		//TMP[countofcolor++] = i;
		GetAns(edge.to);
		if (edge.totalcolor)
		{
			TMP.push_back(i);
			countofcolor++;
			//TMP[countofcolor++] = i;
		}
	}
	res[node][0] = 0;
	res[node][1] = 0;

	long long tmp1 = 1;
	long long tmp0 = 1;
	if (!Color[node])
	{
		tmp1 = 0;
	}
	for (int i = 0; i < countofcolor; i++)
	{

		if (Color[node])
		{
			Edge& edge = G[node][TMP[i]];
			tmp1 *= (res[edge.to][1] + res[edge.to][0]);
			MOD(tmp1);
			tmp0 = 0;
		}
		else
		{
			Edge& edge1 = G[node][TMP[i]];
			tmp0 *= (res[edge1.to][1] + res[edge1.to][0]);
			MOD(tmp0);
			long long tmp3 = 1;
			for (int j = 0; j < countofcolor; j++)
			{
				Edge& edge = G[node][TMP[j]];
				if (i == j)
				{
					tmp3 *= res[edge.to][1];
					MOD(tmp3);
				}
				else
				{
					tmp3 *= (res[edge.to][1] + res[edge.to][0]);
					MOD(tmp3);
				}

			}
			tmp1 += tmp3;

		}

		if (i == countofcolor - 1)
		{
			res[node][0] += tmp0;
			res[node][1] += tmp1;
			MOD(res[node][0]);
			MOD(res[node][1]);
		}

	}
	if (countofcolor == 0)
	{
		res[node][0] = Color[node] ? 0 : 1;
		res[node][1] = Color[node] ? 1 : 0;
	}
}

int main()
{
#ifdef _DEBUG
	freopen("e:\\in.txt", "r", stdin);
#endif // _DEBUG
	int n;
	scanf("%d", &n);
	for (int i = 0; i < n - 1; i++)
	{
		int value;
		scanf("%d", &value);
		AddEdge(i + 1, value);
	}
	for (int i = 0; i < n; i++)
	{
		int value;
		scanf("%d", &value);
		Color[i] = value;
	}
	memset(Visited, 0, sizeof(Visited));
	CountColor(0);
	memset(Visited, 0, sizeof(Visited));
	GetAns(0);
	printf("%I64d\n", res[0][1]);
	return 0;
}