test/yosupo/tree_path_composite_sum.test.cpp

View this file on GitHub
Last update: 2024-01-07 23:25:49+09:00
Problem: https://judge.yosupo.jp/problem/tree_path_composite_sum

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/tree_path_composite_sum"

#include <bits/stdc++.h>

#include "../../math/modint.hpp"
#include "../../tree/rerooting.hpp"
using namespace std;
using ll = long long;

using mint = Modint<998244353>;

struct AddCountMonoid {
    using T = pair<mint, int>;
    static T id() { return {0, 0}; }
    static T op(T a, T b) { return {a.first + b.first, a.second + b.second}; }
};

vector<mint> a;

pair<mint, int> apply_edge(pair<mint, int> x, int, int, pair<mint, mint> e) {
    return {e.first * x.first + e.second * x.second, x.second};
}
pair<mint, int> apply_vertex(pair<mint, int> x, int v) {
    return {x.first + a[v], x.second + 1};
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    a.resize(N);
    for (auto& x : a) cin >> x;
    Rerooting<AddCountMonoid, pair<mint, mint>, apply_edge, apply_vertex>
        rerooting(N);

    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        mint b, c;
        cin >> u >> v >> b >> c;
        rerooting.add_edge(u, v, {b, c});
    }
    auto ans = rerooting.run();
    for (int i = 0; i < N; ++i) {
        cout << ans[i].first << (i < N - 1 ? " " : "\n");
    }
}

#line 1 "test/yosupo/tree_path_composite_sum.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/tree_path_composite_sum"

#include <bits/stdc++.h>

#line 4 "math/modint.hpp"

/**
 * @brief Mod int
 */
template <int m>
class Modint {
    using mint = Modint;
    static_assert(m > 0, "Modulus must be positive");

   public:
    static constexpr int mod() { return m; }

    constexpr Modint(long long y = 0) : x(y >= 0 ? y % m : (y % m + m) % m) {}

    constexpr int val() const { return x; }

    constexpr mint& operator+=(const mint& r) {
        if ((x += r.x) >= m) x -= m;
        return *this;
    }
    constexpr mint& operator-=(const mint& r) {
        if ((x += m - r.x) >= m) x -= m;
        return *this;
    }
    constexpr mint& operator*=(const mint& r) {
        x = static_cast<int>(1LL * x * r.x % m);
        return *this;
    }
    constexpr mint& operator/=(const mint& r) { return *this *= r.inv(); }

    constexpr bool operator==(const mint& r) const { return x == r.x; }

    constexpr mint operator+() const { return *this; }
    constexpr mint operator-() const { return mint(-x); }

    constexpr friend mint operator+(const mint& l, const mint& r) {
        return mint(l) += r;
    }
    constexpr friend mint operator-(const mint& l, const mint& r) {
        return mint(l) -= r;
    }
    constexpr friend mint operator*(const mint& l, const mint& r) {
        return mint(l) *= r;
    }
    constexpr friend mint operator/(const mint& l, const mint& r) {
        return mint(l) /= r;
    }

    constexpr mint inv() const {
        int a = x, b = m, u = 1, v = 0;
        while (b > 0) {
            int t = a / b;
            std::swap(a -= t * b, b);
            std::swap(u -= t * v, v);
        }
        return mint(u);
    }

    constexpr mint pow(long long n) const {
        mint ret(1), mul(x);
        while (n > 0) {
            if (n & 1) ret *= mul;
            mul *= mul;
            n >>= 1;
        }
        return ret;
    }

    friend std::ostream& operator<<(std::ostream& os, const mint& r) {
        return os << r.x;
    }

    friend std::istream& operator>>(std::istream& is, mint& r) {
        long long t;
        is >> t;
        r = mint(t);
        return is;
    }

   private:
    int x;
};
#line 4 "tree/rerooting.hpp"

template <typename M, typename Cost,
          typename M::T (*apply_edge)(typename M::T, int, int, Cost),
          typename M::T (*apply_vertex)(typename M::T, int)>
class Rerooting {
    using T = M::T;

   public:
    explicit Rerooting(int n) : G(n) {}

    void add_edge(int u, int v, Cost c) {
        G[u].emplace_back(v, c);
        G[v].emplace_back(u, c);
    }

    std::vector<T> run() {
        dp_sub.resize(G.size(), M::id());
        dp_all.resize(G.size());
        dfs_sub(0, -1);
        dfs_all(0, -1, M::id());
        return dp_all;
    }

   private:
    std::vector<std::vector<std::pair<int, Cost>>> G;
    std::vector<T> dp_sub, dp_all;

    void dfs_sub(int v, int p) {
        for (auto [c, cost] : G[v]) {
            if (c == p) continue;
            dfs_sub(c, v);
            dp_sub[v] = M::op(dp_sub[v], apply_edge(dp_sub[c], v, c, cost));
        }
        dp_sub[v] = apply_vertex(dp_sub[v], v);
    }

    void dfs_all(int v, int p, const T& val) {
        std::vector<T> ds = {val};
        for (auto [c, cost] : G[v]) {
            if (c == p) continue;
            ds.push_back(apply_edge(dp_sub[c], v, c, cost));
        }
        int n = ds.size();
        std::vector<T> head(n + 1, M::id()), tail(n + 1, M::id());
        for (int i = 0; i < n; ++i) head[i + 1] = M::op(head[i], ds[i]);
        for (int i = n - 1; i >= 0; --i) tail[i] = M::op(ds[i], tail[i + 1]);
        dp_all[v] = apply_vertex(head[n], v);
        int k = 1;
        for (auto [c, cost] : G[v]) {
            if (c == p) continue;
            dfs_all(c, v,
                    apply_edge(apply_vertex(M::op(head[k], tail[k + 1]), v), c,
                               v, cost));
            ++k;
        }
    }
};
#line 7 "test/yosupo/tree_path_composite_sum.test.cpp"
using namespace std;
using ll = long long;

using mint = Modint<998244353>;

struct AddCountMonoid {
    using T = pair<mint, int>;
    static T id() { return {0, 0}; }
    static T op(T a, T b) { return {a.first + b.first, a.second + b.second}; }
};

vector<mint> a;

pair<mint, int> apply_edge(pair<mint, int> x, int, int, pair<mint, mint> e) {
    return {e.first * x.first + e.second * x.second, x.second};
}
pair<mint, int> apply_vertex(pair<mint, int> x, int v) {
    return {x.first + a[v], x.second + 1};
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    a.resize(N);
    for (auto& x : a) cin >> x;
    Rerooting<AddCountMonoid, pair<mint, mint>, apply_edge, apply_vertex>
        rerooting(N);

    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        mint b, c;
        cin >> u >> v >> b >> c;
        rerooting.add_edge(u, v, {b, c});
    }
    auto ans = rerooting.run();
    for (int i = 0; i < N; ++i) {
        cout << ans[i].first << (i < N - 1 ? " " : "\n");
    }
}