結果
問題 | No.2115 Making Forest Easy |
ユーザー |
![]() |
提出日時 | 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 |
ソースコード
#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 : 頂点iの部分木のみを考えた時、頂点iを含む連結成分内のaのmaxがjになるような切り方の数// dp[i][j].second : 頂点iの部分木のみを考えた時、頂点iを含む連結成分内のaのmaxがjになるような切り方全てについての、iを含む連結成分のサイズの和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^imint res = 0;void dfs(int i, int prev) {dp[i] = {{a[i], {1, 1}}}; // 頂点iの1頂点,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]) {// iの部分木外で自由に決められる辺は、iの部分木内の辺と(存在するなら)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;}