code/comp1.cpp at master · ACnoway/code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 1e5 + 5, mod = 998244353, inv2 = (mod + 1) / 2;

int read()
{
    int s = 0, w = 1;
    char ch = getchar();
    while (ch < '0' || ch > '9')
    {
        if (ch == '-')
            w = -w;
        ch = getchar();
    }
    while (ch >= '0' && ch <= '9')
    {
        s = (s << 1) + (s << 3) + (ch ^ '0');
        ch = getchar();
    }
    return s * w;
}
int n, ans, a[N], fa[N], suba[N], du[N], dv[N], f[N], g[N];
vector<int> GA[N];
inline void chksum(int x, int &y)
{
    y += x;
    if (y >= mod)
        y -= mod;
}
inline int Div(int x) { return x %= mod, (x & 1) ? ((x * inv2) % mod) : (x >> 1); }
void dfs(int u, int F)
{
    suba[u] = a[u], fa[u] = F;
    for (int v : GA[u])
    {
        if (v ^ F)
            dfs(v, u), chksum(suba[v], suba[u]);
    }
}
signed main()
{
    freopen("aw.in", "r", stdin);
    freopen("aw.out", "w", stdout);
    n = read();
    for (int i = 1; i <= n; i++)
        f[i] = a[i] = read(), g[i] = (a[i] * a[i]) % mod;
    for (int i = 1; i < n; i++)
        du[i] = read(), GA[du[i]].emplace_back(dv[i] = read());
    dfs(1, 0);
    for (int i = 1; i < n; i++)
    {
        int x = du[i], y = dv[i];
        if (fa[x] == y)
            swap(x, y);

        int t = suba[y], s = (suba[1] - t + mod) % mod;
        int fx = f[x], fy = f[y], gx = g[x], gy = g[y];
        chksum(Div(s * t + (s + mod - t) * fx + mod - gx), ans);
        chksum(Div(s * t + (t + mod - s) * fy + mod - gy), ans);
        f[x] = f[y] = Div(fx + fy), g[x] = g[y] = Div(gx + gy + 2 * fx * fy);
    }
    for (int i = 1; i < n; i++)
        (ans <<= 1) %= mod;
    return cout << (ans % mod + mod) % mod, 0;
}