結果
| 問題 |
No.2949 Product on Tree
|
| コンテスト | |
| ユーザー |
 cureskol
|
| 提出日時 | 2024-10-25 21:53:19 |
| 言語 | C++23 (gcc 13.3.0 + boost 1.87.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 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 3 |
| other | AC * 46 |
ソースコード
#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);
}