【DP】 HDOJ 2974 Counting heaps

dp 공식을 내놓으면 할 수 있어요.

#include <cstdio>
#include <vector>
#include <map>
#include <cstring>
#include <queue>
using namespace std;
typedef long long LL;

#define mp(x, y) make_pair(x, y)
const int maxn = 500005;

queue<int> q;
int fa[maxn];
int in[maxn];
int size[maxn];
bool p[maxn];
int pr[maxn];
int c[maxn];
int n, m, pcnt;

void init()
{
	memset(c, 0, sizeof c);
	memset(in, 0, sizeof in);
}

void Init()
{
	int N = 500005;
	for(int i = 2; i <= N; i++) if(p[i] == 0) {
		for(int j = i + i; j <= N; j += i) p[j] = 1;
	}
	pcnt = 0;
	for(int i = 2; i <= N; i++) if(p[i] == 0) pr[pcnt++] = i;
}

void solve(int x, int flag)
{
	if(p[x] == 0) {
		c[x] += flag;
		return;
	}
	for(int i = 0; (LL)pr[i] * pr[i] <= x; i++) {
		if(x % pr[i] == 0) {
			int cnt = 0;
			while(x % pr[i] == 0) x /= pr[i], cnt++;
			c[pr[i]] += flag * cnt;
		}
		if(p[x] == 0) {
			c[x] += flag;
			return;
		}
	}
}

LL powmod(LL a, LL b, LL mod)
{
	LL res = 1, base = a;
	while(b) {
		if(b % 2) res = res * base % mod;
		base = base * base % mod;
		b /= 2;
	}
	return res;
}

void bfs()
{
	for(int i = 1; i <= n; i++) if(in[i] == 0) q.push(i);
	while(!q.empty()) {
		int u = q.front();
		q.pop();
		size[fa[u]] += size[u];
		in[fa[u]]--;
		if(in[fa[u]] == 0) q.push(fa[u]);
	}
}

void work()
{
	scanf("%d%d", &n, &m);
	for(int i = 1; i <= n; i++) size[i] = 1;
	for(int i = 2; i <= n; i++) {
		scanf("%d", &fa[i]);
		in[fa[i]]++;
	}
	
	bfs();
	for(int i = 2; i < size[1]; i++) solve(i, 1);	
	for(int i = 2; i <= n; i++) solve(size[i], -1);
	LL ans = 1;
	for(int i = 2; i < size[1]; i++) {
		ans = ans * powmod(i, c[i], m) % m;
	}
	printf("%I64d
", ans);
}

int main()
{
//freopen("data", "r", stdin);
	Init();
	int _;
	scanf("%d", &_);
	while(_--) {
		init();
		work();
	}
	
	return 0;
}