UVA 1436 - Counting heaps

題目鏈接

題意：給定一個樹的結構，放1-n數字進去，父親結點值必須小於子節點，問情況有幾種.

思路：f[u]表示以u為子樹的情況，那麼子樹情況為f(v1), f(v2), f(v3)... f(vn).去組成子樹相當於從中選s(v1), s(v2), s(v3) ... s(vn).根據組合數學，情況為f(v1) f(v2) ... f(vn) (s(u) - 1)! \ (s(v1)! s(v2)! ... * s(vn)!)化簡後得到公式：

f(u) = (s(root) - 1)! / s(1)s(2)s(3)...s(n)

由於答案很大要模上m，但是m不是素數，不能直接求逆，只能把分子分母分別分解質因子去消，最後得到都是分子形式在去取模。

代碼：

#include 
#include 
#include 
using namespace std;

const int N = 500005;
int t, n, cnt[N], prime[N], vis[N], pn = 0, f[N], ispri[N];
long long m;

int bfs() {//用dfs必然RE
	queue Q;
	for (int i = 1; i <= n; i++)
		if (!vis[i]) Q.push(i);
 	while (!Q.empty()) {
 		int now = Q.front();
 		Q.pop();
 		if (now == 0) break;
 		cnt[f[now]] += cnt[now];
 		vis[f[now]]--;
 		if (vis[f[now]] == 0)
 			Q.push(f[now]);
  	}
}

void solve(int num, int v) {
	for (int i = 0; i < pn && prime[i] <= num; i++) {
		while (num % prime[i] == 0) {
			cnt[prime[i]] += v;
			num /= prime[i];
  		}
  		if (ispri[num]) {//不加此剪枝必然TLE
  			cnt[num] += v;
     		        break;
    	        }
 	}
}

long long pow_mod(long long x, int k) {
	long long ans = 1;
	while (k) {
		if (k&1) ans = ans * x % m;
		x = x * x % m;
		k >>= 1;
 	}
	return ans;
}

int main() {
	for (int i = 2; i < N; i++) {
		if (vis[i]) continue;
		ispri[i] = 1;
		prime[pn++] = i;
		for (int j = i; j < N; j += i)
			vis[j] = 1;
 	}
	scanf("%d", &t);
	while (t--) {
		scanf("%d%lld", &n, &m);
		for (int i = 1; i <= n; i++)
			cnt[i] = 1;
		f[1] = 0;
		memset(vis, 0, sizeof(vis));
		for (int i = 2; i <= n; i++) {
			scanf("%d", &f[i]);
			vis[f[i]]++;
  		}
  		bfs();
    	memset(vis, 0, sizeof(vis));
  		for (int i = 1; i <= n; i++)
  			vis[cnt[i]]++;
		memset(cnt, 0, sizeof(cnt));
		for (int i = 2; i <= n; i++) {
			solve(i, 1);
			if (vis[i])
   				solve(i, -vis[i]);
		}
		long long ans = 1;
		for (int i = 0; i < pn; i++) {
			if (cnt[prime[i]] == 0) continue;
  			ans = (ans * pow_mod((long long)prime[i], cnt[prime[i]])) % m;
		}
  		printf("%lld\n", ans);
 	}
	return 0;
}