CWOI S C0299 C 樹題

題意

U2FsdGVkX18T278dBirMbWpY+y+yG0mooqC+jJM9ya7VIROpGnZxEEQgrWtQsjbA
DsQyQjhS3U0XUQih8b7ha79KxO8NW6YvusKeAg6XNR6eZVrk1pBAxeso8UWZFvXA
Fsqf4cTh4GXYqARxnS8SOrcskt8a8oI9py3W7AvFuLTm6zZlOi+nTnJF46QUzgit
wl3CFGCaa5H/v84J/UotFQaYmWyac8Dh594MZ3A9zXbqCQsWuN/jLNhczAUjkWm7
LWnxtUVMLa+hI1D6kmj9ac4Wxmu3PDn8qwZcn4oF8GShezmziu+GQUTRMsek2x7+
nDqtdTql95xEgdcNPEz/GVuQ10/SK/4WPgqnIGRGHyBSx17XklDUW8vcLezpjpuq
sx0EUHy76atPk3pufi/Ta8CFzRJzoC0/sYfEG7+ig4SauMm9rxgYpV85DWnBYf/x

解析

先考慮求解數量。令 $f(i, 0/1/2)$ 表示以 $i$ 為根節點的子樹，$0$ 表示 $i$ 處選了一個完整的四元組，或者沒有選，$1$ 表示 $i$ 處選了 $1$ 個節點，$2$ 表示 $i$ 處選了 $2$ 個節點的鏈。

對於 $f(i, 1)$，容易得到 $f(i, 1) = \prod_{v} f(v, 0)$。對於 $f(i, 2)$ 要求兒子中恰有一個選擇是選 $1$ 個節點，即：$f(i, 2) = \sum_{v} f(v, 1) \prod_{v' \ne v} f(v, 0)$。而 $f(i, 0)$ 則需要恰有一個選 $1$ 且恰有一個選 $2$，或者全部不選，即： $f(i, 0) = \prod_{v} f(v, 0) + \sum_{v \ne v'} f(v, 1)f(v', 2) \prod_{v'' \ne v \land v'' \ne v'} f(v'', 0)$。為了方便，可以用一個輔助的 dp 陣列 $g(0/1, 0/1)$ 表示當前是否選了 1 和是否選了 2 的方案數。

然後考慮求解和。一個簡便的方法是，沿用上面的 dp 方法，將每個元素當成一個二元組 $(c, s)$，表示數量和和。考慮在二元組上定義 $+$ 和 $\times$。由定義，容易得到：

$$\begin{aligned} (c_{1}, s_{1}) + (c_{2}, s_{2}) &= (c_{1} + c_{2}, s_{1} + s_{2}) \\ (c_{1}, s_{1}) \times (c_{2}, s_{2}) &= (c_{1}c_{2}, s_{1}c_{2} + s_{2}c_{1}) \\ \end{aligned}$$

然後根據上面的運算律計算即可。

實現

struct Solve
{
	struct node_t
	{
		mint cnt, sum;
		node_t() : cnt(0), sum(0) {}
		node_t(mint cnt, mint sum) : cnt(cnt), sum(sum) {}
		node_t operator+=(node_t a) 
		{
			cnt += a.cnt;
			sum += a.sum;
			return *this;
		};
		node_t operator*=(node_t a) 
		{
			sum = sum * a.cnt + a.sum * cnt;
			cnt *= a.cnt;
			return *this;
		};
		node_t operator*=(mint x)
		{
			cnt *= x;
			sum *= x;
			return *this;
		}
		node_t operator/=(mint x)
		{
			cnt /= x;
			sum /= x;
			return *this;
		}
		node_t operator+(node_t a) const { return node_t{*this} += a; }
		node_t operator*(node_t a) const { return node_t{*this} *= a; }
		node_t operator*(mint x) const { return node_t{*this} *= x; }
		node_t operator/(mint x) const { return node_t{*this} /= x; }
	};

	int n;
	std::vector<std::vector<int>> adj;
	std::vector<int> fa;
	std::vector<std::array<node_t, 3>> f;

	Solve(int n) : n(n), adj(n), fa(n), f(n)
	{
		for (int i = 0; i < n - 1; i++) {
			int u, v;
			std::cin >> u >> v;
			u--;
			v--;
			adj[u].emplace_back(v);
			adj[v].emplace_back(u);
		}
	}

	void dfs_init(int u)
	{
		for (auto v : adj[u]) {
			if (v == fa[u]) continue;
			fa[v] = u;
			dfs_init(v);
		}
	}

	void dfs_dp(int u)
	{
		f[u][0] = f[u][2] = node_t{0, 0};
		f[u][1] = node_t{1, 0};
		for (auto v : adj[u]) {
			if (v == fa[u]) continue;
			dfs_dp(v);
		}
		for (auto v : adj[u]) {
			if (v == fa[u]) continue;
			f[u][1] *= f[v][0];
		}
		std::array<std::array<node_t, 2>, 2> g;
		g[0][0] = node_t{1, 0};
		g[0][1] = g[1][0] = g[1][1] = node_t{0, 0};
		for (auto v : adj[u]) {
			if (v == fa[u]) continue;
			std::array<std::array<node_t, 2>, 2> h;
			h[0][0] += g[0][0] * f[v][0];
			h[0][1] += g[0][0] * f[v][1] + g[0][1] * f[v][0];
			h[1][0] += g[0][0] * f[v][2] + g[1][0] * f[v][0];
			h[1][1] += g[0][1] * f[v][2] + g[1][0] * f[v][1] + g[1][1] * f[v][0];
			g = std::move(h);
		}
		f[u][2] = g[0][1];
		f[u][0] = f[u][1] + g[1][1] * node_t{1, 1};
	}

	void solve()
	{
		dfs_init(0);
		dfs_dp(0);
	}
};

int main()
{
	set_io("tree");

	int n, type;
	std::cin >> n >> type;
	Solve s(n);
	s.solve();
	auto res = s.f[0][0];
	std::cout << res.cnt.val() << std::endl;
	if (type) std::cout << res.sum.val() << std::endl;
}