結果

問題 No.2360 Path to Integer
コンテスト
ユーザー 遭難者
提出日時 2024-03-22 08:16:59
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 83 ms / 2,500 ms
コード長 5,222 bytes
コンパイル時間 8,269 ms
コンパイル使用メモリ 349,996 KB
実行使用メモリ 25,756 KB
最終ジャッジ日時 2024-09-30 10:30:08
合計ジャッジ時間 10,323 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 15
権限があれば一括ダウンロードができます

ソースコード

diff # 

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#else
#pragma GCC target("avx2")
#pragma GCC optimize("Ofast,unroll-loops")
#endif
#include <bits/stdc++.h>
using namespace std;
#if __has_include(<atcoder/all>)
#include <atcoder/all>
using mint = atcoder::modint998244353;
istream& operator>>(istream& is, mint& a) {
	int t; is >> t; a = t;
	return is;
}
ostream& operator<<(ostream& os, mint a) {
	return os << a.val();
}
#endif
#define rep(i, n) for (int i = 0; i < n; i++)
#define per(i, n) for (int i = n - 1; i >= 0; i--)
#define ALL(a) a.begin(), a.end()
#define vec vector
#undef long
#define long long long
template<typename T>
ostream& operator<<(ostream& os, vector<T> a) {
	const int n = a.size();
	rep(i, n) os << a[i] << " \n"[i + 1 == n];
	return os;
}
template<typename T, size_t n>
ostream& operator<<(ostream& os, array<T, n> a) {
	rep(i, n) os << a[i] << " \n"[i + 1 == n];
	return os;
}
template<typename T>
istream& operator>>(istream& is, vector<T>& a) {
	for (T& i : a) is >> i;
	return is;
}
template<typename T>
bool chmin(T& x, T y) {
	if (x > y) { x = y; return true; }
	return false;
}
template<typename T>
bool chmax(T& x, T y) {
	if (x < y) { x = y; return true; }
	return false;
}
template<typename T>
void operator++(vector<T>& a) {
	for (T& i : a) ++i;
}
template<typename T>
void operator--(vector<T>& a) {
	for (T& i : a) --i;
}
template<typename T>
void operator++(vector<T>& a, int) {
	for (T& i : a) i++;
}
template<typename T>
void operator--(vector<T>& a, int) {
	for (T& i : a) i--;
}
#undef endl
#define endl '\n'
template <class E, class V, E(*merge)(E, E), E(*e)(), E(*put_edge)(V, int), V(*put_vertex)(E, int)>
struct RerootingDP {
	struct edge {
		int to, idx, xdi;
	};
	RerootingDP(int n_ = 0) : n(n_), inner_edge_id(0) {
		es.resize(2 * n - 2);
		start.resize(2 * n - 2);
		if (n == 1) es_build();
	}
	void add_edge(int u, int v, int idx, int xdi) {
		start[inner_edge_id] = u;
		es[inner_edge_id] = { v,idx,xdi };
		inner_edge_id++;
		start[inner_edge_id] = v;
		es[inner_edge_id] = { u,xdi,idx };
		inner_edge_id++;
		if (inner_edge_id == 2 * n - 2) {
			es_build();
		}
	}
	vector<V> build(int root_ = 0) {
		root = root_;
		vector<V> subdp(n); subdp[0] = put_vertex(e(), 0);
		outs.resize(n);
		vector<int> geta(n + 1, 0);
		for (int i = 0; i < n; i++) geta[i + 1] = start[i + 1] - start[i] - 1;
		geta[root + 1]++;
		for (int i = 0; i < n; i++) geta[i + 1] += geta[i];
		auto dfs = [&](auto sfs, int v, int f) -> void {
			E val = e();
			for (int i = start[v]; i < start[v + 1]; i++) {
				if (es[i].to == f) {
					swap(es[start[v + 1] - 1], es[i]);
				}
				if (es[i].to == f) continue;
				sfs(sfs, es[i].to, v);
				E nval = put_edge(subdp[es[i].to], es[i].idx);
				outs[geta[v]++] = nval;
				val = merge(val, nval);
			}
			subdp[v] = put_vertex(val, v);
			};
		dfs(dfs, root, -1);
		return subdp;
	}
	vector<V> reroot() {
		vector<E> reverse_edge(n);
		reverse_edge[root] = e();
		vector<V> answers(n);
		auto dfs = [&](auto sfs, int v) -> void {
			int le = outs_start(v);
			int ri = outs_start(v + 1);
			int siz = ri - le;
			vector<E> rui(siz + 1);
			rui[siz] = e();
			for (int i = siz - 1; i >= 0; i--) {
				rui[i] = merge(outs[le + i], rui[i + 1]);
			}
			answers[v] = put_vertex(merge(rui[0], reverse_edge[v]), v);
			E lui = e();
			for (int i = 0; i < siz; i++) {
				V rdp = put_vertex(merge(merge(lui, rui[i + 1]), reverse_edge[v]), v);
				reverse_edge[es[start[v] + i].to] = put_edge(rdp, es[start[v] + i].xdi);
				lui = merge(lui, outs[le + i]);
				sfs(sfs, es[start[v] + i].to);
			}
			};
		dfs(dfs, root);
		return answers;
	}
private:
	int n, root, inner_edge_id;
	vector<E> outs;
	vector<edge> es;
	vector<int> start;
	int outs_start(int v) {
		int res = start[v] - v;
		if (root < v) res++;
		return res;
	}
	void es_build() {
		vector<edge> nes(2 * n - 2);
		vector<int> nstart(n + 2, 0);
		for (int i = 0; i < 2 * n - 2; i++) nstart[start[i] + 2]++;
		for (int i = 0; i < n; i++) nstart[i + 1] += nstart[i];
		for (int i = 0; i < 2 * n - 2; i++) nes[nstart[start[i] + 1]++] = es[i];
		swap(es, nes);
		swap(start, nstart);
	}
};
constexpr int INF = 1e5 + 2;
mint ten[12];
mint a_val[INF], a_len[INF];
struct mint_int {
	mint val;
	int size;
	mint_int() : val(0), size(0) {}
	mint_int(mint val, int size) : val(val), size(size) {}
};
mint_int merge(mint_int x, mint_int y) { return mint_int(x.val + y.val, x.size + y.size); }
mint_int e() { return mint_int(0, 0); }
mint_int put_edge(mint_int x, int) { return x; }
mint_int put_vertex(mint_int x, int idx) { return mint_int((x.size + 1) * a_val[idx] + x.val * a_len[idx], x.size + 1); }
void solve() {
	ten[0] = 1;
	for (int i = 1; i < 12; i++) ten[i] = 10 * ten[i - 1];
	int n; cin >> n;
	rep(i, n) {
		string s; cin >> s;
		a_val[i] = stol(s);
		a_len[i] = ten[s.size()];
	}
	RerootingDP<mint_int, mint_int, merge, e, put_edge, put_vertex> cd(n);
	rep(i, n - 1) {
		int u, v; cin >> u >> v;
		u--, v--;
		cd.add_edge(u, v, v, u);
	}
	cd.build();
	mint ans = 0;
	for (auto [val, idx] : cd.reroot())
		ans += val;
	cout << ans << endl;
}
int main() {
	// srand((unsigned)time(NULL));
	cin.tie(nullptr);
	ios::sync_with_stdio(false);
	// cout << fixed << setprecision(40);
	solve();
	return 0;
}
0