模拟赛题讲解[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：

输出样例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;
}