結果

問題 No.2949 Product on Tree
ユーザー noya2noya2
提出日時 2024-11-02 17:51:21
言語 C++23
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 543 ms / 2,000 ms
コード長 3,620 bytes
コンパイル時間 4,245 ms
コンパイル使用メモリ 269,664 KB
実行使用メモリ 72,700 KB
最終ジャッジ日時 2024-11-02 17:51:54
合計ジャッジ時間 29,147 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge3
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
6,820 KB
testcase_01 AC 2 ms
6,816 KB
testcase_02 AC 2 ms
6,816 KB
testcase_03 AC 460 ms
44,676 KB
testcase_04 AC 477 ms
43,788 KB
testcase_05 AC 462 ms
44,920 KB
testcase_06 AC 458 ms
45,016 KB
testcase_07 AC 453 ms
44,032 KB
testcase_08 AC 467 ms
45,532 KB
testcase_09 AC 499 ms
46,040 KB
testcase_10 AC 472 ms
46,708 KB
testcase_11 AC 472 ms
48,108 KB
testcase_12 AC 496 ms
52,468 KB
testcase_13 AC 474 ms
53,180 KB
testcase_14 AC 485 ms
56,112 KB
testcase_15 AC 466 ms
60,428 KB
testcase_16 AC 481 ms
57,384 KB
testcase_17 AC 502 ms
58,672 KB
testcase_18 AC 482 ms
58,484 KB
testcase_19 AC 495 ms
62,416 KB
testcase_20 AC 483 ms
61,696 KB
testcase_21 AC 490 ms
62,380 KB
testcase_22 AC 496 ms
63,648 KB
testcase_23 AC 455 ms
45,452 KB
testcase_24 AC 487 ms
45,608 KB
testcase_25 AC 460 ms
45,756 KB
testcase_26 AC 462 ms
45,856 KB
testcase_27 AC 477 ms
45,936 KB
testcase_28 AC 489 ms
45,992 KB
testcase_29 AC 484 ms
46,508 KB
testcase_30 AC 477 ms
47,592 KB
testcase_31 AC 490 ms
49,720 KB
testcase_32 AC 515 ms
50,592 KB
testcase_33 AC 486 ms
57,084 KB
testcase_34 AC 510 ms
65,244 KB
testcase_35 AC 497 ms
60,032 KB
testcase_36 AC 488 ms
68,508 KB
testcase_37 AC 543 ms
69,140 KB
testcase_38 AC 503 ms
63,676 KB
testcase_39 AC 525 ms
64,560 KB
testcase_40 AC 502 ms
71,880 KB
testcase_41 AC 535 ms
67,224 KB
testcase_42 AC 500 ms
72,700 KB
testcase_43 AC 174 ms
33,552 KB
testcase_44 AC 175 ms
34,080 KB
testcase_45 AC 226 ms
42,820 KB
testcase_46 AC 201 ms
38,740 KB
testcase_47 AC 148 ms
29,220 KB
testcase_48 AC 208 ms
39,824 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include<bits/stdc++.h>
#include<atcoder/modint>

// g[from] contains outgoing edges (to, edgeid(from, to))
// (E, op, e) is commutative monoid
// ~edgeid(from, to) == edgeid(to, from)
// return calculator of dp(r, v)
template<class V, class E>
auto rerootingdp(auto op, E e, auto put_edge, auto put_vertex, const std::vector<std::vector<std::pair<int, int>>> &g, int root = 0){
    struct dp {
        // dp(r, v) : root is r, dp of subtree v
        // ans[v]  = dp(v, v)
        // from[v] = dp(v, par(v))
        // to[v]   = dp(par(v), v)
        // from[root] and to[root] is undefined
        std::vector<V> ans, from, to;
        std::vector<int> down, up;
        std::vector<std::vector<int>> childs;
        bool is_in_subtree(int r, int v){
            return down[r] < down[v] && up[v] <= up[r];
        }
        V operator()(int r, int v){
            if (r == v) return ans[v];
            if (!is_in_subtree(v, r)) return to[v];
            int le = 0, ri = childs[v].size();
            while (ri - le > 1){
                int md = (le + ri) / 2;
                if (down[childs[v][md]] <= down[r]){
                    le = md;
                }
                else {
                    ri = md;
                }
            }
            return from[childs[v][le]];
        }
    };
    int n = g.size();
    std::vector<E> from(n, e), to(n, e);
    std::vector<V> dp_to(n);
    std::vector<std::vector<int>> childs(n);
    std::vector<int> down(n), up(n);
    int t = 0;
    auto dfs = [&](auto sfs, int v, int f) -> void {
        down[v] = t++;
        childs[v].reserve(g[v].size());
        for (auto [c, eid] : g[v]){
            if (c == f) continue;
            childs[v].emplace_back(c);
            sfs(sfs, c, v);
            dp_to[c] = put_vertex(to[c], c);
            to[c] = put_edge(dp_to[c], eid);
            to[v] = op(to[v], to[c]);
        }
        up[v] = t;
    };
    dfs(dfs, root, -1);
    std::vector<V> dp_ans(n), dp_from(n);
    std::vector<E> inner(n);
    auto bfs = [&](auto sfs, int v, int f) -> void {
        int sz = g[v].size();
        inner[sz] = e;
        int i = sz-1;
        for (auto [c, eid] : g[v] | std::views::reverse){
            if (c == f) continue;
            inner[i] = op(inner[i+1], to[c]);
            i--;
        }
        dp_ans[v] = put_vertex(op(inner[++i], from[v]), v);
        E rui = e;
        for (auto [c, eid] : g[v]){
            if (c == f) continue;
            dp_from[c] = put_vertex(op(op(rui, inner[++i]), from[v]), v);
            from[c] = put_edge(dp_from[c], ~eid);
            rui = op(rui, to[c]);
        }
        for (auto [c, eid] : g[v]){
            if (c == f) continue;
            sfs(sfs, c, v);
        }
    };
    bfs(bfs, root, -1);
    return dp{dp_ans, dp_from, dp_to, down, up, childs};
}

using mint = atcoder::modint998244353;

int main(){
    int n; std::cin >> n;
    std::vector<mint> a(n);
    for (auto &e : a){
        int x; std::cin >> x;
        e = x;
    }
    std::vector<std::vector<std::pair<int,int>>> g(n);
    for (int i = 0; i < n-1; i++){
        int u, v; std::cin >> u >> v; u--, v--;
        g[u].emplace_back(v,i);
        g[v].emplace_back(u,~i);
    }
    auto op = [&](mint l, mint r){
        return l + r;
    };
    auto pute = [&](mint v, int){
        return v;
    };
    auto putv = [&](mint e, int v){
        return e * a[v] + a[v];
    };
    auto dp = rerootingdp<mint,mint>(op,mint(0),pute,putv,g);
    mint ans = 0;
    for (int i = 0; i < n; i++){
        ans += dp(i,i);
        ans -= a[i];
    }
    ans /= 2;
    std::cout << ans.val() << '\n';
}
0