結果

問題 No.2115 Making Forest Easy
ユーザー QCFiumQCFium
提出日時 2022-08-22 21:14:15
言語 C++14
(gcc 13.2.0 + boost 1.83.0)
結果
AC  
実行時間 833 ms / 2,000 ms
コード長 4,847 bytes
コンパイル時間 2,060 ms
コンパイル使用メモリ 191,580 KB
実行使用メモリ 258,980 KB
最終ジャッジ日時 2023-09-04 23:28:47
合計ジャッジ時間 18,919 ms
ジャッジサーバーID
（参考情報） 		judge11 / judge14
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
4,380 KB
testcase_01 AC 1 ms
4,380 KB
testcase_02 AC 818 ms
253,120 KB
testcase_03 AC 427 ms
4,544 KB
testcase_04 AC 711 ms
184,024 KB
testcase_05 AC 538 ms
126,896 KB
testcase_06 AC 819 ms
255,796 KB
testcase_07 AC 49 ms
11,284 KB
testcase_08 AC 826 ms
255,632 KB
testcase_09 AC 21 ms
6,740 KB
testcase_10 AC 833 ms
253,288 KB
testcase_11 AC 430 ms
4,596 KB
testcase_12 AC 20 ms
6,728 KB
testcase_13 AC 381 ms
4,388 KB
testcase_14 AC 21 ms
6,824 KB
testcase_15 AC 47 ms
10,780 KB
testcase_16 AC 420 ms
4,608 KB
testcase_17 AC 11 ms
4,376 KB
testcase_18 AC 52 ms
20,044 KB
testcase_19 AC 354 ms
108,564 KB
testcase_20 AC 41 ms
10,620 KB
testcase_21 AC 11 ms
6,192 KB
testcase_22 AC 40 ms
4,380 KB
testcase_23 AC 557 ms
157,448 KB
testcase_24 AC 430 ms
4,796 KB
testcase_25 AC 193 ms
47,176 KB
testcase_26 AC 21 ms
6,728 KB
testcase_27 AC 180 ms
4,376 KB
testcase_28 AC 50 ms
11,680 KB
testcase_29 AC 25 ms
11,596 KB
testcase_30 AC 48 ms
11,376 KB
testcase_31 AC 710 ms
181,232 KB
testcase_32 AC 833 ms
255,540 KB
testcase_33 AC 830 ms
258,980 KB
testcase_34 AC 16 ms
4,376 KB
testcase_35 AC 87 ms
32,140 KB
testcase_36 AC 823 ms
254,396 KB
testcase_37 AC 167 ms
4,380 KB
testcase_38 AC 66 ms
4,380 KB
testcase_39 AC 238 ms
77,076 KB
testcase_40 AC 44 ms
10,676 KB
testcase_41 AC 21 ms
6,756 KB
testcase_42 AC 436 ms
4,528 KB
testcase_43 AC 25 ms
7,832 KB
testcase_44 AC 64 ms
4,376 KB
testcase_45 AC 35 ms
9,240 KB
testcase_46 AC 422 ms
4,596 KB
testcase_47 AC 18 ms
6,584 KB
testcase_48 AC 697 ms
171,532 KB
testcase_49 AC 734 ms
190,252 KB
testcase_50 AC 232 ms
4,376 KB
testcase_51 AC 119 ms
4,380 KB
権限があれば一括ダウンロードができます

ソースコード

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  : 頂点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^i
mint 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;
}
0