#1022546 (C++23) No.2949 Product on Tree

提出ソース

結果

問題	No.2949 Product on Tree
ユーザー	cureskol
提出日時	2024-10-25 21:53:19
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	141 ms / 2,000 ms
コード長	14,882 bytes
コンパイル時間	6,014 ms
コンパイル使用メモリ	316,020 KB
実行使用メモリ	25,304 KB
最終ジャッジ日時	2024-10-25 21:53:50
合計ジャッジ時間	14,731 ms
ジャッジサーバーID （参考情報）	judge5 / judge1

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,812 KB
testcase_01	AC	2 ms 6,820 KB
testcase_02	AC	2 ms 6,816 KB
testcase_03	AC	125 ms 19,592 KB
testcase_04	AC	123 ms 19,212 KB
testcase_05	AC	129 ms 19,536 KB
testcase_06	AC	128 ms 19,640 KB
testcase_07	AC	119 ms 19,284 KB
testcase_08	AC	121 ms 19,824 KB
testcase_09	AC	127 ms 19,860 KB
testcase_10	AC	124 ms 19,924 KB
testcase_11	AC	129 ms 20,188 KB
testcase_12	AC	134 ms 21,052 KB
testcase_13	AC	127 ms 21,144 KB
testcase_14	AC	121 ms 21,576 KB
testcase_15	AC	123 ms 22,448 KB
testcase_16	AC	124 ms 21,844 KB
testcase_17	AC	123 ms 22,156 KB
testcase_18	AC	135 ms 22,208 KB
testcase_19	AC	128 ms 22,904 KB
testcase_20	AC	132 ms 22,856 KB
testcase_21	AC	132 ms 23,136 KB
testcase_22	AC	123 ms 22,940 KB
testcase_23	AC	124 ms 20,012 KB
testcase_24	AC	125 ms 19,900 KB
testcase_25	AC	128 ms 19,964 KB
testcase_26	AC	129 ms 19,896 KB
testcase_27	AC	126 ms 19,868 KB
testcase_28	AC	126 ms 20,024 KB
testcase_29	AC	122 ms 20,020 KB
testcase_30	AC	126 ms 20,172 KB
testcase_31	AC	126 ms 20,688 KB
testcase_32	AC	141 ms 20,676 KB
testcase_33	AC	133 ms 22,084 KB
testcase_34	AC	133 ms 23,876 KB
testcase_35	AC	131 ms 22,768 KB
testcase_36	AC	138 ms 24,508 KB
testcase_37	AC	141 ms 24,548 KB
testcase_38	AC	138 ms 23,484 KB
testcase_39	AC	135 ms 23,608 KB
testcase_40	AC	135 ms 25,012 KB
testcase_41	AC	135 ms 24,152 KB
testcase_42	AC	131 ms 25,304 KB
testcase_43	AC	58 ms 15,764 KB
testcase_44	AC	58 ms 15,824 KB
testcase_45	AC	75 ms 19,392 KB
testcase_46	AC	68 ms 17,888 KB
testcase_47	AC	50 ms 13,972 KB
testcase_48	AC	70 ms 18,160 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#pragma region template
#ifdef USE_PRECOMPILE
#include "pch.h"
#else
#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using ld = long double;
using vi = std::vector<int>;
using vll = std::vector<ll>;
template <typename T> using PA = std::pair<T, T>;
using pi = PA<int>;
using pll = PA<ll>;

#define overload2(a, b, c, ...) c
#define overload3(a, b, c, d, ...) d
#define overload4(a, b, c, d, e, ...) e
#define overload5(a, b, c, d, e, f, ...) f

#define TYPE1(T) template <typename T>
#define TYPE2(T, U) template <typename T, typename U>
#define TYPE(...) overload2(__VA_ARGS__, TYPE2, TYPE1)(__VA_ARGS__)
#define TYPES1(T) template <typename... T>
#define TYPES2(H, T) template <typename H, typename... T>
#define TYPES(...) overload2(__VA_ARGS__, TYPES2, TYPES1)(__VA_ARGS__)

#define REP4(i, s, n, d) for (int i = (s); i < (n); i += (d))
#define REP3(i, s, n) REP4(i, s, n, 1)
#define REP2(i, n) REP3(i, 0, n)
#define REP1(n) REP2(tomato, n)
#define REP(...) overload4(__VA_ARGS__, REP4, REP3, REP2, REP1)(__VA_ARGS__)

#define RREP4(i, n, s, d) for (int i = (n) - 1; i >= (s); i -= (d))
#define RREP3(i, n, s) RREP4(i, n, s, 1)
#define RREP2(i, n) RREP3(i, n, 0)
#define RREP1(n) RREP2(tomato, n)
#define RREP(...)                                                              \
    overload4(__VA_ARGS__, RREP4, RREP3, RREP2, RREP1)(__VA_ARGS__)

#define FOR4(a, b, c, d, v) for (auto [a, b, c, d] : v)
#define FOR3(a, b, c, v) for (auto [a, b, c] : v)
#define FOR2(a, b, v) for (auto [a, b] : v)
#define FOR1(a, v) for (auto a : v)
#define FOR(...) overload5(__VA_ARGS__, FOR4, FOR3, FOR2, FOR1)(__VA_ARGS__)

#define AFOR4(a, b, c, d, v) for (auto &[a, b, c, d] : v)
#define AFOR3(a, b, c, v) for (auto &[a, b, c] : v)
#define AFOR2(a, b, v) for (auto &[a, b] : v)
#define AFOR1(a, v) for (auto &a : v)
#define AFOR(...)                                                              \
    overload5(__VA_ARGS__, AFOR4, AFOR3, AFOR2, AFOR1)(__VA_ARGS__)

#define CFOR4(a, b, c, d, v) for (const auto &[a, b, c, d] : v)
#define CFOR3(a, b, c, v) for (const auto &[a, b, c] : v)
#define CFOR2(a, b, v) for (const auto &[a, b] : v)
#define CFOR1(a, v) for (const auto &a : v)
#define CFOR(...)                                                              \
    overload5(__VA_ARGS__, CFOR4, CFOR3, CFOR2, CFOR1)(__VA_ARGS__)

#define RFOR4(a, b, c, d, v) for (auto [a, b, c, d] : v | std::views::reverse)
#define RFOR3(a, b, c, v) for (auto [a, b, c] : v | std::views::reverse)
#define RFOR2(a, b, v) for (auto [a, b] : v | std::views::reverse)
#define RFOR1(a, v) for (auto a : v | std::views::reverse)
#define RFOR(...)                                                              \
    overload5(__VA_ARGS__, RFOR4, RFOR3, RFOR2, RFOR1)(__VA_ARGS__)

#define RAFOR4(a, b, c, d, v) for (auto &[a, b, c, d] : v | std::views::reverse)
#define RAFOR3(a, b, c, v) for (auto &[a, b, c] : v | std::views::reverse)
#define RAFOR2(a, b, v) for (auto &[a, b] : v | std::views::reverse)
#define RAFOR1(a, v) for (auto &a : v | std::views::reverse)
#define RAFOR(...)                                                             \
    overload5(__VA_ARGS__, RAFOR4, RAFOR3, RAFOR2, RAFOR1)(__VA_ARGS__)

#define RCFOR4(a, b, c, d, v)                                                  \
    for (const auto &[a, b, c, d] : v | std::views::reverse)
#define RCFOR3(a, b, c, v) for (const auto &[a, b, c] : v | std::views::reverse)
#define RCFOR2(a, b, v) for (const auto &[a, b] : v | std::views::reverse)
#define RCFOR1(a, v) for (const auto &a : v | std::views::reverse)
#define RCFOR(...)                                                             \
    overload5(__VA_ARGS__, RCFOR4, RCFOR3, RCFOR2, RCFOR1)(__VA_ARGS__)

#define ALL(v) v.begin(), v.end()
#define RALL(v) v.rbegin(), v.rend()
#define SORT(v) std::ranges::sort(v)
#define RSORT(v) sort(RALL(v))
#define REVERSE(v) std::ranges::reverse(v)
#define UNIQUE(v) SORT(v), v.erase(unique(ALL(v)), v.end())
#define SZ(v) int(v.size())
#define MAX(v) std::ranges::max(v)
#define MIN(v) std::ranges::min(v)
template <typename T>
std::conditional_t<std::is_same_v<T, int>, ll, T> SUM(const std::vector<T> &v) {
    return std::accumulate(
        ALL(v), std::conditional_t<std::is_same_v<T, int>, ll, T>(0));
}

TYPES(T) void input(T &...a) { (std::cin >> ... >> a); }
#define DECLARE(T, ...)                                                        \
    T __VA_ARGS__;                                                             \
    input(__VA_ARGS__);
#define INT(...) DECLARE(int, __VA_ARGS__)
#define STR(...) DECLARE(std::string, __VA_ARGS__)
#define LL(...) DECLARE(long long, __VA_ARGS__)
#define CHR(...) DECLARE(char, __VA_ARGS__)
#define DBL(...) DECLARE(double, __VA_ARGS__)
#define VI(n, v)                                                               \
    vi v(n);                                                                   \
    std::cin >> v;
#define VLL(n, v)                                                              \
    vll v(n);                                                                  \
    std::cin >> v;
#define VS(n, s)                                                               \
    std::vector<std::string> s(n);                                             \
    std::cin >> s;

TYPE(T, S) std::istream &operator>>(std::istream &is, std::pair<T, S> &p) {
    is >> p.first >> p.second;
    return is;
}
TYPE(T, S)
std::ostream &operator<<(std::ostream &os, const std::pair<T, S> &p) {
    os << p.first << " " << p.second;
    return os;
}
TYPE(T) istream &operator>>(std::istream &is, std::vector<T> &v) {
    for (auto &a : v)
        std::cin >> a;
    return is;
}
TYPE(T) std::ostream &operator<<(std::ostream &os, const std::vector<T> &v) {
    if (&os == &std::cerr)
        os << "[";
    REP (i, v.size()) {
        os << v[i];
        if (i + 1 < v.size())
            os << (&os == &std::cerr ? "," : " ");
    }
    if (&os == &std::cerr)
        os << "]";
    return os;
}

#ifdef __DEBUG
#include <debug>
#else
#define debug(...) void(0)
#endif

void print() { cout << '\n'; }
TYPES(T, Ts) void print(const T &a, const Ts &...b) {
    std::cout << a;
    (std::cout << ... << (std::cout << ' ', b));
    std::cout << '\n';
}

TYPE(T) using pq = std::priority_queue<T>;
TYPE(T) using pqg = std::priority_queue<T, std::vector<T>, std::greater<T>>;
TYPE(T) T pick(std::queue<T> &que) {
    assert(que.size());
    T a = que.front();
    que.pop();
    return a;
}
TYPE(T) T pick(pq<T> &que) {
    assert(que.size());
    T a = que.top();
    que.pop();
    return a;
}
TYPE(T) T pick(pqg<T> &que) {
    assert(que.size());
    T a = que.top();
    que.pop();
    return a;
}
TYPE(T) T pick(std::stack<T> &sta) {
    assert(sta.size());
    T a = sta.top();
    sta.pop();
    return a;
}
TYPE(T) void clear(T &v) { v = decltype(v)(); }

std::string YES(bool f = true) { return (f ? "YES" : "NO"); }
std::string Yes(bool f = true) { return (f ? "Yes" : "No"); }
std::string yes(bool f = true) { return (f ? "yes" : "no"); }

constexpr int INF = 1e9 + 7;
constexpr ll LINF = ll(1e18) + 7;
constexpr ld EPS = 1e-10;

vi iota(int n) {
    vi a(n);
    std::iota(ALL(a), 0);
    return a;
}
TYPE(T) void add(std::vector<T> &v, T a = 1) {
    AFOR (p, v)
        p += a;
}
TYPE(T) T rev(T a) {
    REVERSE(a);
    return a;
}
TYPE(T) void fin(T a) {
    std::cout << a << '\n';
    exit(0);
}
TYPE(T) bool chmax(T &a, T b) { return (a < b && (a = b, true)); }
TYPE(T) bool chmin(T &a, T b) { return (a > b && (a = b, true)); }
TYPES(T, Ns) auto make_vector(T x, int n, Ns... ns) {
    if constexpr (sizeof...(ns) == 0)
        return std::vector<T>(n, x);
    else
        return std::vector(n, make_vector<T>(x, ns...));
}
bool in(const ll S, const int a) { return (S >> a) & 1; }
int digit(char c) { return (c >= '0' and c <= '9' ? c - '0' : -1); }

void signal_hander(int signal) { std::exit(EXIT_FAILURE); }
#endif
#pragma endregion template

#include <atcoder/convolution>
#include <atcoder/modint>
using namespace atcoder;
using mint = modint998244353;
namespace atcoder {
ostream &operator<<(ostream &os, mint a) {
    os << a.val();
    return os;
}
istream &operator>>(istream &is, mint &a) {
    long long b;
    is >> b;
    a = b;
    return is;
}
} // namespace atcoder
template <typename TREE, typename Data> class ReRooting {
    int n;
    TREE T;

  public:
    ReRooting(const TREE &T) : T(T), n(T.n) {}

    template <typename F1, typename F2>
    std::vector<Data> build(const F1 &score, const F2 &merge,
                            const Data &unit) {
        if (!~T.root)
            T.build(0);
        std::vector<Data> dp1(n, unit), dp2(n);
        for (int v : T.DFS)
            for (const auto &e : T.son(v)) {
                dp2[e.to] = score(dp1[e.to], e);
                merge(dp1[v], dp2[e.to]);
            }
        std::vector<Data> top(n, unit), res(n);
        for (int v : T.BFS) {
            Data now = (T.root == v ? unit : score(top[v], T.parent(v)));
            for (int to : T.son(v)) {
                top[to] = now;
                merge(now, dp2[to]);
            }
            res[v] = now;
            now = unit;
            for (int i = T.son(v).size() - 1; i >= 0; i--) {
                int to = T.son(v)[i];
                merge(top[to], now);
                merge(now, dp2[to]);
            }
        }
        return res;
    }
};

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

struct Edge {
    int from, to;
    Edge() = default;
    Edge(int from, int to) : from(from), to(to) {}
    operator int() const { return to; }
};

struct Graph {
    int n;
    using edge_type = Edge;
    std::vector<edge_type> edges;

  protected:
    std::vector<int> in_deg;
    bool prepared;
    class OutgoingEdges {
        Graph *g;
        int l, r;

      public:
        OutgoingEdges(Graph *g, int l, int r) : g(g), l(l), r(r) {}
        edge_type *begin() { return &(g->edges[l]); }
        edge_type *end() { return &(g->edges[r]); }
        edge_type &operator[](int i) { return g->edges[l + i]; }
        int size() const { return r - l; }
    };
    class ConstOutgoingEdges {
        const Graph *g;
        int l, r;

      public:
        ConstOutgoingEdges(const Graph *g, int l, int r) : g(g), l(l), r(r) {}
        const edge_type *begin() const { return &(g->edges[l]); }
        const edge_type *end() const { return &(g->edges[r]); }
        const edge_type &operator[](int i) const { return g->edges[l + i]; }
        int size() const { return r - l; }
    };

  public:
    OutgoingEdges operator[](int v) {
        assert(prepared);
        return {this, in_deg[v], in_deg[v + 1]};
    }
    const ConstOutgoingEdges operator[](int v) const {
        assert(prepared);
        return {this, in_deg[v], in_deg[v + 1]};
    }

    bool is_prepared() const { return prepared; }

    Graph() : n(0), in_deg(1, 0), prepared(false) {}
    Graph(int n) : n(n), in_deg(n + 1, 0), prepared(false) {}
    Graph(int n, int m, bool directed = false, int indexed = 1)
        : n(n), in_deg(n + 1, 0), prepared(false) {
        scan(m, directed, indexed);
    }

    void resize(int n) { n = n; }

    void add_arc(int from, int to) {
        assert(!prepared);
        assert(0 <= from and from < n and 0 <= to and to < n);
        edges.emplace_back(from, to);
        in_deg[from + 1]++;
    }
    void add_edge(int u, int v) {
        add_arc(u, v);
        add_arc(v, u);
    }
    void add_arc(const edge_type &e) { add_arc(e.from, e.to); }
    void add_edge(const edge_type &e) { add_edge(e.from, e.to); }

    void scan(int m, bool directed = false, int indexed = 1) {
        edges.reserve(directed ? m : 2 * m);
        while (m--) {
            int u, v;
            std::cin >> u >> v;
            u -= indexed;
            v -= indexed;
            if (directed)
                add_arc(u, v);
            else
                add_edge(u, v);
        }
        build();
    }

    void build() {
        assert(!prepared);
        prepared = true;
        for (int v = 0; v < n; v++)
            in_deg[v + 1] += in_deg[v];
        std::vector<edge_type> new_edges(in_deg.back());
        auto counter = in_deg;
        for (auto &&e : edges)
            new_edges[counter[e.from]++] = e;
        edges = new_edges;
    }

    void graph_debug() const {
#ifndef __LOCAL
        return;
#endif
        assert(prepared);
        for (int from = 0; from < n; from++) {
            std::cerr << from << ";";
            for (int i = in_deg[from]; i < in_deg[from + 1]; i++)
                std::cerr << edges[i].to << " ";
            std::cerr << "\n";
        }
    }
};
struct Tree : Graph {
    using Graph::Graph;
    Tree() = default;
    int root = -1;
    std::vector<int> DFS, BFS, depth;

    void scan_root(int indexed = 1) {
        for (int i = 1; i < n; i++) {
            int p;
            std::cin >> p;
            add_edge(p - indexed, i);
        }
        build();
    }
    void scan(int indexed = 1) {
        Graph::scan(n - 1, false, indexed);
        build();
    }

    edge_type &parent(int v) {
        assert(~root and root != v);
        return (*this)[v][0];
    }
    const edge_type &parent(int v) const {
        assert(~root and root != v);
        return (*this)[v][0];
    }

    OutgoingEdges son(int v) {
        assert(~root);
        if (v == root)
            return {this, in_deg[v], in_deg[v + 1]};
        return {this, in_deg[v] + 1, in_deg[v + 1]};
    }

  private:
    void dfs(int v, int pre = -1) {
        for (auto &e : (*this)[v]) {
            if (e.to == pre)
                std::swap((*this)[v][0], e);
            else {
                depth[e.to] = depth[v] + 1;
                dfs(e.to, v);
            }
        }
        DFS.push_back(v);
    }

  public:
    void build(int r = 0) {
        if (!is_prepared())
            Graph::build();
        if (~root) {
            assert(r == root);
            return;
        }
        root = r;
        depth = std::vector<int>(n, 0);
        DFS.reserve(n);
        BFS.reserve(n);
        dfs(root);
        std::queue<int> que;
        que.push(root);
        while (que.size()) {
            int p = que.front();
            que.pop();
            BFS.push_back(p);
            for (const auto &e : son(p))
                que.push(e.to);
        }
    }
};

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::signal(SIGABRT, signal_hander);

    INT(n);
    VI(n, v);
    Tree t(n);
    t.scan();

    ReRooting<Tree, mint> R(t);

    auto score = [&](mint x, auto e) { return (x + 1) * v[e.to]; };
    auto merge = [&](mint &x, mint y) { x += y; };

    auto d = R.build(score, merge, mint::raw(0));
    debug(d);
    mint ans = 0;
    REP (i, n)
        ans += d[i] * v[i];
    fin(ans / 2);
}

yukicoder

結果

テストケース

ソースコード