TC SRM421 R1 Div1 TeamManagement（树形 dp + 计数 dp）

Description

有 $n$ 个球员，有的球员对俱乐部忠诚，有的球员对俱乐部不忠诚。有的球员是朋友，朋友的关系是一棵无根树，具有对称性但不具有传递性，每个球员最多有三个朋友。想购买一个球员，需要满足这个球员是忠诚的或这个球员的朋友被购买。俱乐部想购买 $k$ 个球员，求每个球员被买到的概率。

$1\le k\le n\le 50$

Solution

有一个不同寻常的性质：每个球员最多有三个朋友。所以选择一个度为以的点 $rt$ 为根，每个点最多有两个儿子。，

进行树形 dp。令 $f_{x,cnt,typ}$ 为以 $x$ 为根的子树中，选择 $cnt$ 个点的方案数。定义根据 $typ$ 的值略有改动：

• $typ=0$： 无改动。
• $typ=1$： $x$ 与子树外的忠诚点相连。
• $typ=2$：$x$ 与子树内的忠诚点相连。

考虑转移。

$x$ 为叶子节点，$cnt\le 1$。

当 $cnt=0$ 时：$typ=0,1,2$ 对应的 $1,1,0$。当 $cnt=1$ 时：$x$ 为忠诚的为 $1$，$x$ 不忠诚时：$typ=0,1,2$ 对应的 $0,1,0$。

$x$ 为非叶子节点 且 $x$ 有两个儿子 $y_1$ 和 $y_2$。

• 不选择 $x$。

  如果以 $y_1$ 为根的子树选择了 $i$ 个点，那么 $y_2$ 为根的子树需要选择 $cnt–i$ 个点。根据乘法原理为 $f_{y_1,i,0}\times f_{y_2,cnt-i,0}$。对 $i\in [0,cnt]$ 的答案求和。

• 选择 $x$ 。

  • $typ=1$ 或 $x$ 是忠诚的。对于 $y_1$ 和 $y_2$ 它们连到了某个子树外忠诚点（可以是 $x$），所以答案是 $f_{y_1,i,1}\times f_{y_2,cnt-i-1,1}$ 的总和。

  • $typ\ne 1$ 且 $x$ 不忠诚。想要选择 $x$，必须让 $x$ 与忠诚点相连，所以在 $y_1$ 或 $y_2$ 的中，要有至少一个 $typ=2$。看上去答案是：$f_{y_1,i,2}\times f_{y_2,cnt-i-1,1}+f_{y_1,i,1}\times f_{y_2,cnt-i-1,2}$。但是 $typ=2$ 的答案当 $y_1$ 和 $y_2$ 两个子树中都有忠诚节点是会被重复计算，所以答案为：$f_{y_1,i,2}\times f_{y_2,cnt-i-1,1}+f_{y_1,i,1}\times f_{y_2,cnt-i-1,2}-f_{y_1,i,2}\times f_{y_2,cnt-i-1,2}$。

至于怎么计算概率。可以禁止一个点被选，求出方案数，用 $1-$ 禁止与不禁止的方案数之比求概率。

Code

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 50 + 5; 
class TeamManagement {	
	private:
		int n, k;
		int a[N][N], deg[N], canTake[N], isLoyal[N]; 
		ll f[N][N][3];
	public:
		void dfs(int x, int fa) {
			// find children
			vector <int> cd;
			for (int y = 0; y < n; y++) if (a[x][y] && y != fa) {
				dfs(y, x); cd.push_back(y); 
			}
			int sz = cd.size();
			if (sz == 0) { // leaf vertex
				if (isLoyal[x]) {
					f[x][0][0] = 1, f[x][1][0] = (canTake[x] ? 1 : 0);
					f[x][0][1] = 1, f[x][1][1] = (canTake[x] ? 1 : 0);
					f[x][0][2] = 0, f[x][1][2] = (canTake[x] ? 1 : 0);
				} else {
					f[x][0][0] = 1, f[x][1][0] = 0;
       	 			f[x][0][1] = 1, f[x][1][1] = (canTake[x] ? 1 : 0);
        			f[x][0][2] = 0, f[x][1][2] = 0;
				}
				return ;
			}
			//Don't take x
			for (int i = 0; i <= k; i++) {
				if (sz == 1) f[x][i][0] += f[cd[0]][i][0], f[x][i][1] += f[cd[0]][i][0];
				else {
					for (int j = 0; j <= i; j++) {
						f[x][i][0] += f[cd[0]][j][0] * f[cd[1]][i - j][0],
						f[x][i][1] += f[cd[0]][j][0] * f[cd[1]][i - j][0];
					}
				}
			}
			// Take x 
			if (!canTake[x]) return ;
			for (int i = 1; i <= k; i++)
				for (int typ = 0; typ <= 2; typ++) 
					if (typ == 1 || isLoyal[x]) {
						if (sz == 1) f[x][i][typ] += f[cd[0]][i - 1][1];
						else for (int j = 0; j < i; j++) f[x][i][typ] += f[cd[0]][j][1] * f[cd[1]][i - j - 1][1];
					} else {
						if (sz == 1) f[x][i][typ] += f[cd[0]][i - 1][2];
						else for (int j = 0; j < i; j++) f[x][i][typ] += f[cd[0]][j][1] * f[cd[1]][i - j - 1][2] 
						+ f[cd[0]][j][2] * f[cd[1]][i - j - 1][1] - f[cd[0]][j][2] * f[cd[1]][i - j - 1][2]; 
					}
		}
		vector <double> getDistribution (int _n, int _k, vector <string> friends, string loyal) {
			n = _n, k = _k;
			for (int i = 0; i < n; i++) isLoyal[i] = (loyal[i] == 'Y');
			for (int i = 0; i < n - 1; i++) {
				istringstream sin(friends[i]);
				int x, y; sin >> x >> y;
				a[x][y] = a[y][x] = 1; 
				deg[x]++, deg[y]++;
			}
			int rt = 0;
			while (deg[rt] > 1) rt++; //Take a leaf as  a root 
			memset(canTake, 1, sizeof(canTake)); memset(f, 0, sizeof(f));
			vector <double> ans; 
			dfs(rt, -1); ll tot = f[rt][k][0];
			for (int i = 0; i < n; i++) {
				memset(f, 0, sizeof(f));
				canTake[i] = 0; dfs(rt, -1); canTake[i] = 1;
				ans.push_back(1.0 - (f[rt][k][0] * 1.0 / tot));
			}
			return ans;
		}
}Solve;