模拟赛题讲解[26]
来自 s_r_f 2023-08-02 noi.ac #3194
题目描述:
小 D 有一棵 \(n\) 个点的树 \(T\),点有点权 \(a_i\),边有边权 \(w_i\)。
记 \(d(x,y)\) 为 \(x\) 到 \(y\) 简单路径上的边权和,\(f(x) = \sum\limits_{y=1}^n d(x,y)^c \times a_y\),对 \(i = 1,2,\cdots n\) 求 \(f(i)\) 的值。
答案对 998244353 取模。
输入格式:
第一行,两个整数 \(n\) 和 \(c\),用空格隔开 。
接下来 \(1\) 行 \(n\) 个整数 \(a_1,a_2 \cdots a_{n}\) ,用空格隔开。
接下来 \(n-1\) 行,每行三个整数 \(u,v,w\) 表示一条连接点 \(u,v\),边权为 \(w\) 的边
输出格式:
输出一行。
一行 \(n\) 个数 \(f(1),f(2),\cdots,f(n)\) ,对 998244353 取模的结果。
输入样例1:
3 1
1 1 1
1 2 10
2 3 100
输出样例1:
120 110 210
数据范围:
子任务 1(20pts) : \(n \leq 1000\)
子任务 2(20pts) \(:\) \(T\) 是一条链,即保证存在一个排列 \((p_1,\cdots p_n)\) 满足边集恰好是 \(\{(p_1,p_2),... (p_{n-1},p_n)\}\)
子任务 3(20pts) \(:\) \(c \leq 1\)
子任务 4(40pts) : \(2 \leq n \leq 200000,0 \leq c \leq 3,a_i, w \leq 10^9\)
题解:
总觉得好像在哪做过,但是忘了。
设 \(f(u,c)\) 表示 \(u\) 所在子树的第 \(c\) 层答案,转移如下:( \(a\uparrow b\) 表示 a += b
)
\[
\begin{aligned}
&f_{u,0} \uparrow f_{v,0}
\\[6pt]&f_{u,1} \uparrow f_{v,1} + wf_{v,0}
\\[6pt]&f_{u,2} \uparrow f_{v,2} + 2w\times f_{v,1} + w^2f_{v,0}
\\[6pt]&f_{u,3} \uparrow f_{v,3} + 3wf_{v,2} + 3w^2f_{v,1} +
w^3f_{v,0}
\end{aligned}
\] 然后做一个换根dp就好啦。
时间复杂度 \(\mathcal{O}(n)\)
代码:
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
const int mod = 998244353;
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
void add(int &x,int y) { (x += y % mod) >= mod ? x -= mod : (x < 0 ? x += mod : 0); }
#define N ((int)(2e5 + 15))
struct Edge { int u,v,w,next; } e[N << 1];
int n,c,pos,head[N],a[N],f[N][4];
void addEdge(int u,int v,int w) { e[++pos] = {u,v,w,head[u]}; head[u] = pos; }
void dfs1(int u,int fa = 0)
{
f[u][0] = a[u];
for(int i = head[u]; i; i = e[i].next)
{
int v = e[i].v, w = e[i].w; if(v == fa) continue;
/*
dp[x][0] += dp[y][0];
dp[x][1] += dp[y][1] + dp[y][0] * v;
dp[x][2] += dp[y][2] + 2 * dp[y][1] * v + dp[y][0] * v^2;
dp[x][3] += dp[y][3] + 3 * dp[y][2] * v + 3 * dp[y][1] * v^2 + dp[y][0] * v^3;
*/
dfs1(v,u);
add(f[u][0], f[v][0]);
add(f[u][1], f[v][1] + w * f[v][0]);
add(f[u][2], f[v][2] + 2 * w * f[v][1]);
add(f[u][2], w * w % mod * f[v][0] % mod);
add(f[u][3], f[v][3] + 3 * w * f[v][2]);
add(f[u][3], 3 * w * w % mod * f[v][1] % mod);
add(f[u][3], w * w % mod * w % mod * f[v][0] % mod);
}
}
void solve(int u[4], int v[4], int w, int au, int av)
{
int t0 = u[0], t1 = u[1], t2 = u[2], t3 = u[3];
add(u[0], -v[0]);
add(u[1], -v[1] - w * v[0]);
add(u[2], -v[2] - 2 * w * v[1]);
add(u[2], -w * w % mod * v[0] % mod);
add(u[3], -v[3] - 3 * w * v[2]);
add(u[3], -3 * w * w % mod * v[1] % mod);
add(u[3], -w * w % mod * w % mod * v[0] % mod);
add(v[0], u[0]);
add(v[1], u[1] + w * u[0]);
add(v[2], u[2] + 2 * w * u[1]);
add(v[2], w * w % mod * u[0] % mod);
add(v[3], u[3] + 3 * w * u[2]);
add(v[3], 3 * w * w % mod * u[1] % mod);
add(v[3], w * w % mod * w % mod * u[0] % mod);
u[0] = t0, u[1] = t1, u[2] = t2, u[3] = t3;
}
void dfs2(int u,int fa = 0)
{
for(int i = head[u]; i; i = e[i].next) {
if(e[i].v != fa) {
solve(f[u], f[e[i].v], e[i].w, a[u], a[e[i].v]); dfs2(e[i].v, u);
}
}
}
signed main()
{
ios::sync_with_stdio(0);
cin.tie(0);cout.tie(0);
// freopen("check.in","r",stdin);
// freopen("check.out","w",stdout);
cin >> n >> c;
for(int i = 1; i <= n; i++) cin >> a[i];
for(int i = 1,u,v,w; i < n; i++) {
cin >> u >> v >> w; addEdge(u,v,w); addEdge(v,u,w);
}
dfs1(1); dfs2(1);
for(int i = 1; i <= n; i++) cout << f[i][c] << " \n"[i == n];
return 0;
}