P7846 「dWoi R2」Arcade hall / 街机厅 题解

P7846 「dWoi R2」Arcade hall / 街机厅（2）

题意简述：有一棵 $ n $ 个点的树，第 $ i $ 个点有点权 $ w_i $，第 $ j $ 条边有边权 $ t_j $；每一条边 $ (u_j,v_j,t_j) $ 两边点的点权有如下要求：

• \(t_j=0\)，\(w_{u_j} \neq w_{v_j}\)；
• \(t_j=1\)，没有要求；
• \(t_j=2\)，\(w_{u_j}=w_{v_j}\)；
• 对任意点权都要有 \(w_i \in [1,R]\)；

求当 $ w_i $ 作为序列时，一共有多少种本质不同的序列 $ w_i $ 以及 $ w_i $ 的和的最小值。

思路：首先，根据题目中在树上的问题以及答案只与父子节点之间的权值有关，我们考虑树形 DP。观察到 \(R\le 100\)，\(N\le 10^5\) 的数据范围，结合到本题需要用到节点权值进行计算答案的特点，我们可以在节点维度后增加一个点权维度。设 $ f_{u,x} $ 表示节点 $ u $ 的权值 $ w_u=x $ 时，它的子树下的序列 $ w_i $ 数量；$ g_{u,x} $ 表示节点 $ u $ 的权值 $ w_u=x $ 时，它的子树下的 $ w_i $ 和的最小值。则我们可以进行 DFS。对于每条边 $ (u,v,t) $：

• 若 $ t=0 $：则 $ w_v\neq x $，所以 \(f_{u,x}\times \sum_{y\in [1,R]\&y\neq x}f_{v,y}\to f_{u,x}\)，\(g_{u,x}+\min_{y\in[1,R]\&y\neq x}g_{v,y}\to g_{u,x}\)；
• 若 $ t=1 $：则无限制，所以 \(f_{u,x}\times \sum_{y\in[1,R]}f_{v,y}\to f_{u,x}\)，\(g_{u,x}+\min_{y\in[1,R]}g_{v,y}\to g_{u,x}\)；
• 若 $ t=2 $：则 $ w_v=x $，所以 \(f_{u,x}\times f_{v,x}\to f_{u,x}\)，\(g_{u,x}+g_{v,x}\to g_{u,x}\)。

注意：$ f_{u,x} $ 需要初始化为 \(1\)，\(g_{u,x}\) 需要初始化为 \(x\)。

算一下时间复杂度，遍历 $ n $ 个节点，再对每个节点遍历一遍 $ R $ 计算，若 $ t=0/1 $ 还需要遍历一遍 $ R $，总时间复杂度为 $ \mathcal O(nR^2) $，达到了 $ 10^9 $，会 TLE。于是我们考虑优化。

我们发现，时间主要浪费在上面 $ t=0/1 $ 时的遍历 $ y $ 部分了，我们可以对 \(f_{u,x}\) 计算 \(f'_{u}=\sum_{i=1}^{R}f_{u,i}\)，对 \(g_{u,x}\) 计算 \(g^{p}_{u,x}=\min_{i=1}^{x}g_{u,i}\)，\(g^s_{u,x}=\min_{i=x}^Rg_{u,i}\)。则上面遍历 $ y $ 的部分就可以换成：

• 当 $ t=0 $ 时，\(f_{u,x}\times (f'_{v}-f_{v,x})\to f_{u,x}\)，\(g_{u,x}+\min\{g^p_{v,x-1},g^s_{v,x+1}\to g_{u,x}\}\)；
• 当 $ t=1 $ 时，\(f_{u,x}\times f'_v\to f_{u,x}\)，\(g_{u,x}+g^p_{v,R}\to g_{u,x}\)。

这样时间复杂度就优化到了 $ \mathcal O(nR) $，可以通过本题。

代码：

#include<bits/stdc++.h>

#define vec vector
#define pub push_back
#define pob pop_back
#define debug(x) cerr << #x << " = " << x << endl
#define lowbit(x) (x) & (-x)
#define IOS ios::sync_with_stdio(0), cin.tie(0), cout.tie(0)

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const double EPS = 1e-8;
const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3f;
const int N = 1e5 + 5, R = 105;
const ll MOD = 1e9 + 7;

int n, r;
ll f[N][R], g[N][R], f2[N], mgpre[N][R], mgsuf[N][R];
vector<pii> G[N];

void dfs(int x, int fa) {
	for(int i = 1; i <= r; i++) {
		f[x][i] = 1, g[x][i] = i;
	}
	for(pii e : G[x]) {
		int y = e.first, t = e.second;
		if(y == fa) continue;
		dfs(y, x);
		for(int i = 1; i <= r; i++) {
			if(t == 0) {
				f[x][i] = (f[x][i] * (f2[y] - f[y][i]) % MOD + MOD) % MOD;
				g[x][i] = g[x][i] + min(mgpre[y][i - 1], mgsuf[y][i + 1]);
			} else if(t == 1) {
				f[x][i] = f[x][i] * f2[y] % MOD;
				g[x][i] = g[x][i] + mgsuf[y][1];
			} else {
				f[x][i] = f[x][i] * f[y][i] % MOD;
				g[x][i] = g[x][i] + g[y][i];
			}
		}
	}
	for(int i = r; i >= 1; i--) mgsuf[x][i] = min(mgsuf[x][i + 1], g[x][i]);
	for(int i = 1; i <= r; i++) {
		f2[x] = (f2[x] + f[x][i]) % MOD;
		mgpre[x][i] = min(mgpre[x][i - 1], g[x][i]);
		// cerr << "f[" << x << "][" << i << "]=" << f[x][i] << ", g[" << x << "][" << i << \
			"]=" << g[x][i] << endl;
	}
}

void solve() {
	cin >> n >> r;
	for(int i = 1; i < n; i++) {
		int u, v, t; cin >> u >> v >> t;
		G[u].pub({v, t}), G[v].pub({u, t});
	}
	memset(g, 0x3f, sizeof g);
	memset(mgpre, 0x3f, sizeof mgpre);
	memset(mgsuf, 0x3f, sizeof mgsuf);
	dfs(1, 0);
	cout << f2[1] << " " << (f2[1] == 0 ? 0 : mgpre[1][r]);
}

int main() {
	IOS;
	// freopen("a.in", "r", stdin);
	// freopen("a.out", "w", stdout);
	int T = 1;
	// cin >> T;
	while(T--)
		solve();
	return 0;
}