結果

問題 No.2949 Product on Tree
ユーザー InTheBloom
提出日時 2024-11-04 18:19:53
言語 D
(dmd 2.109.1)
結果
AC  
実行時間 414 ms / 2,000 ms
コード長 1,094 bytes
コンパイル時間 2,550 ms
コンパイル使用メモリ 174,312 KB
実行使用メモリ 39,588 KB
最終ジャッジ日時 2024-11-04 18:20:17
合計ジャッジ時間 21,843 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 46
権限があれば一括ダウンロードができます

ソースコード

diff # 

import std;

void main () {
    int N = readln.chomp.to!int;
    auto A = readln.split.to!(int[]);

    const long MOD = 998244353;

    auto graph = new int[][](N, 0);
    foreach (i; 0..N - 1) {
        int U, V; readln.read(U, V);
        U--, V--;
        graph[U] ~= V;
        graph[V] ~= U;
    }

    long ans = 0;
    auto dp = new long[](N);
    // dp[i] := iを根とする部分木に属するすべての頂点jに対して、f(i, j)の総和

    void dfs (int pos, int par) {
        long s = 0;
        foreach (to; graph[pos]) {
            if (to == par) continue;
            dfs(to, pos);

            ans += s * A[pos] % MOD * dp[to] % MOD;
            s += dp[to];
            s %= MOD;
        }
        dp[pos] = (s + 1) * A[pos] % MOD;
        ans += dp[pos] - A[pos];
        ans %= MOD;
        if (ans < 0) ans += MOD;
    }

    dfs(0, -1);

    writeln(ans);
}

void read (T...) (string S, ref T args) {
    import std.conv : to;
    import std.array : split;
    auto buf = S.split;
    foreach (i, ref arg; args) {
        arg = buf[i].to!(typeof(arg));
    }
}
0