結果

問題 No.2115 Making Forest Easy
ユーザー QCFium
提出日時 2022-08-22 21:14:15
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 824 ms / 2,000 ms
コード長 4,847 bytes
コンパイル時間 2,078 ms
コンパイル使用メモリ 192,648 KB
実行使用メモリ 259,200 KB
最終ジャッジ日時 2024-06-22 20:43:31
合計ジャッジ時間 18,311 ms
ジャッジサーバーID
(参考情報)
judge2 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 50
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

#include <bits/stdc++.h>
int ri() {
int n;
scanf("%d", &n);
return n;
}
template<int mod>
struct ModInt{
int x;
ModInt () : x(0) {}
ModInt (int64_t x) : x(x >= 0 ? x % mod : (mod - -x % mod) % mod) {}
ModInt &operator += (const ModInt &p){
if ((x += p.x) >= mod) x -= mod;
return *this;
}
ModInt &operator -= (const ModInt &p) {
if ((x += mod - p.x) >= mod) x -= mod;
return *this;
}
ModInt &operator *= (const ModInt &p) {
x = (int64_t) x * p.x % mod;
return *this;
}
ModInt &operator /= (const ModInt &p) {
*this *= p.inverse();
return *this;
}
ModInt &operator ^= (int64_t p) {
ModInt res = 1;
for (; p; p >>= 1) {
if (p & 1) res *= *this;
*this *= *this;
}
return *this = res;
}
ModInt operator - () const { return ModInt(-x); }
ModInt operator + (const ModInt &p) const { return ModInt(*this) += p; }
ModInt operator - (const ModInt &p) const { return ModInt(*this) -= p; }
ModInt operator * (const ModInt &p) const { return ModInt(*this) *= p; }
ModInt operator / (const ModInt &p) const { return ModInt(*this) /= p; }
ModInt operator ^ (int64_t p) const { return ModInt(*this) ^= p; }
bool operator == (const ModInt &p) const { return x == p.x; }
bool operator != (const ModInt &p) const { return x != p.x; }
explicit operator int() const { return x; }
ModInt &operator = (const int p) { x = p; return *this;}
ModInt inverse() const {
int a = x, b = mod, u = 1, v = 0, t;
while (b > 0) {
t = a / b;
a -= t * b;
std::swap(a, b);
u -= t * v;
std::swap(u, v);
}
return ModInt(u);
}
friend std::ostream & operator << (std::ostream &stream, const ModInt<mod> &p) {
return stream << p.x;
}
friend std::istream & operator >> (std::istream &stream, ModInt<mod> &a) {
int64_t x;
stream >> x;
a = ModInt<mod>(x);
return stream;
}
};
typedef ModInt<998244353> mint;
using dp_t = std::map<int, std::pair<mint, mint> >;
int n;
std::vector<std::vector<int> > hen;
std::vector<int> a;
// dp[i][j].first : iiamaxj
// dp[i][j].second : iiamaxji
std::vector<dp_t> dp;
dp_t merge(const dp_t &lhs, const dp_t &rhs) {
dp_t res;
for (auto i : lhs) {
int key0 = i.first;
auto val0 = i.second;
for (auto j : rhs) {
int key1 = j.first;
auto val1 = j.second;
res[std::max(key0, key1)].first += val0.first * val1.first;
res[std::max(key0, key1)].second += val0.first * val1.second + val0.second * val1.first;
}
}
return res;
}
dp_t merge_fast(const dp_t &lhs, const dp_t &rhs) {
dp_t res;
{
auto itr = rhs.begin();
mint sum0 = 0, sum1 = 0;
for (auto i : lhs) {
int key = i.first;
auto val = i.second;
while (itr != rhs.end() && itr->first < key) {
sum0 += itr->second.first;
sum1 += itr->second.second;
itr++;
}
res[key].first += val.first * sum0;
res[key].second += val.first * sum1 + val.second * sum0;
}
}
{
auto itr = lhs.begin();
mint sum0 = 0, sum1 = 0;
for (auto i : rhs) {
int key = i.first;
auto val = i.second;
while (itr != lhs.end() && itr->first <= key) {
sum0 += itr->second.first;
sum1 += itr->second.second;
itr++;
}
res[key].first += val.first * sum0;
res[key].second += val.first * sum1 + val.second * sum0;
}
}
return res;
}
std::vector<int> subtree_size;
std::vector<mint> power2; // power2[i] : 2^i
mint res = 0;
void dfs(int i, int prev) {
dp[i] = {{a[i], {1, 1}}}; // i1,0
for (auto j : hen[i]) if (j != prev) {
dfs(j, i);
subtree_size[i] += subtree_size[j];
// i - j :
// i - j : 2^(subtree_size[j] - 1)0, max-INF
dp[j][-1] = {power2[subtree_size[j] - 1], 0};
// dp[i] = merge(dp[i], dp[j]);
dp[i] = merge_fast(dp[i], dp[j]);
}
// i
for (auto j : dp[i]) {
// ii()i
res += j.second.second * j.first * power2[n - subtree_size[i] - (prev != -1)];
}
}
int main() {
n = ri();
a.resize(n);
for (auto &i : a) i = ri();
hen.resize(n);
for (int i = 1; i < n; i++) {
int x = ri() - 1;
int y = ri() - 1;
hen[x].push_back(y);
hen[y].push_back(x);
}
dp.resize(n);
subtree_size.resize(n, 1);
power2.resize(n + 1, 1);
for (int i = 1; i <= n; i++) power2[i] = power2[i - 1] + power2[i - 1];
dfs(0, -1);
std::cout << res << std::endl;
return 0;
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0