結果

問題 No.2949 Product on Tree
コンテスト
ユーザー zawakasu
提出日時 2024-10-25 22:05:06
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 296 ms / 2,000 ms
コード長 1,942 bytes
コンパイル時間 2,638 ms
コンパイル使用メモリ 84,852 KB
最終ジャッジ日時 2025-02-24 23:26:56
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 46
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <iostream>
#include <iomanip>
#include <cassert>
#include <vector>

const long long MOD{998244353};

long long add(long long L, long long R) {
    L += R;
    if (L >= MOD) L -= MOD;
    return L;
}

long long inv_plus(long long v) {
    return (v == 0 ? 0LL : MOD - v);
}

long long sub(long long L, long long R) {
    return add(L, inv_plus(R));
}

long long mult(long long L, long long R) {
    return (L * R) % MOD;
}

long long div2(long long v) {
    static long long inv{499122177};
    assert(mult(inv, 2) == 1);
    return mult(v, inv);
}

int main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);

    int N;
    std::cin >> N;
    std::vector<int> A(N);
    for (auto& a : A) std::cin >> a;
    std::vector<std::vector<int>> g(N);
    for (int i{} ; i < N - 1 ; i++) {
        int u, v;
        std::cin >> u >> v;
        u--; v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    std::vector<long long> dp(N);
    auto dfs{[&](auto dfs, int v, int p) -> long long {
        for (auto x : g[v]) if (x != p) {
            dp[v] = add(dp[v], mult(dfs(dfs, x, v), A[v]));
        }
        dp[v] = add(dp[v], A[v]);
        return dp[v];
    }};
    dfs(dfs, 0, -1);
    long long ans{};
    auto dfs2{[&](auto dfs, int v, int p, long long propa) -> void {
        std::vector<long long> ch(g[v].size());
        long long sum{};
        for (int i{} ; i < (int)g[v].size() ; i++) {
            ch[i] = g[v][i] == p ? propa : dp[g[v][i]];
            ch[i] = mult(ch[i], A[v]);
            sum = add(sum, ch[i]);
        } 
        sum = add(sum, A[v]);
        ans = add(ans, sum);
        for (int i{} ; i < (int)g[v].size() ; i++) {
            sum = sub(sum, ch[i]);
            if (g[v][i] != p) dfs(dfs, g[v][i], v, sum);
            sum = add(sum, ch[i]);
        } 
    }};
    dfs2(dfs2, 0, -1, 0LL);
    for (auto a : A) ans = sub(ans, a);
    ans = div2(ans);
    std::cout << ans << '\n';
}
0