結果

問題 No.2360 Path to Integer
ユーザー V_Melville
提出日時 2025-05-02 23:44:29
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 242 ms / 2,500 ms
コード長 1,992 bytes
コンパイル時間 6,046 ms
コンパイル使用メモリ 333,244 KB
実行使用メモリ 50,432 KB
最終ジャッジ日時 2025-05-02 23:44:39
合計ジャッジ時間 9,252 ms
ジャッジサーバーID
(参考情報)
judge5 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 15
権限があれば一括ダウンロードができます

ソースコード

diff #

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace atcoder;
#define rep(i, n) for (int i = 0; i < (n); ++i)

using namespace std;
using ll = long long;
using mint = modint998244353;

vector<mint> ten;
struct Rerooting {
  struct DP {
  	mint x; int w;
    DP(mint x=0, int w=0): x(x), w(w) {}
    DP operator+(const DP& a) const {
      return DP(x+a.x, w+a.w);
    }
    DP addRoot(int c, ll d) const {
      string s = to_string(d);
      int w2 = s.size();
      return DP(x*ten[w2] + mint(d)*(w+c), w+c);
    }
  };
  
  int n;
  vector<ll> a;
  vector<vector<int>> to, cost;
  vector<vector<DP>> dp;
  vector<DP> ans;
  Rerooting(int n=0): n(n),a(n),to(n),cost(n),dp(n),ans(n) {}
  void addEdge(int a, int b, int c) {
    to[a].push_back(b); cost[a].push_back(c);
    to[b].push_back(a); cost[b].push_back(c);
  }
  void init() {
    dfs(0);
    bfs(0);
  }

  DP dfs(int v, int p=-1) {
    DP dpSum;
    dp[v] = vector<DP>(to[v].size());
    rep(i,to[v].size()) {
      int u = to[v][i];
      if (u == p) continue;
      dp[v][i] = dfs(u,v).addRoot(cost[v][i], a[u]);
      dpSum = dpSum + dp[v][i];
    }
    return dpSum;
  }
  void bfs(int v, const DP& dpP=DP(), int p=-1) {
    int deg = to[v].size();
    rep(i,deg) if (to[v][i] == p) dp[v][i] = dpP;

    vector<DP> dpSumL(deg+1);
    rep(i,deg) dpSumL[i+1] = dpSumL[i] + dp[v][i];
    vector<DP> dpSumR(deg+1);
    for (int i = deg-1; i >= 0; --i)
      dpSumR[i] = dpSumR[i+1] + dp[v][i];
    ans[v] = dpSumL[deg].addRoot(1, a[v]);

    rep(i,deg) {
      int u = to[v][i];
      if (u == p) continue;
      DP d = dpSumL[i] + dpSumR[i+1];
      bfs(u, d.addRoot(cost[v][i], a[v]), v);
    }
  }
};

int main() {
	int n;
	cin >> n;

    ten.resize(11, 1);
    rep(i, 10) ten[i+1] = ten[i]*10;

	Rerooting g(n);
	rep(i, n) cin >> g.a[i];
	rep(i, n-1) {
		int a, b;
		cin >> a >> b;
		--a; --b;
		g.addEdge(a, b, 1);
	}
	g.init();
	
	mint ans;
	rep(i, n) ans += g.ans[i].x;
	cout << ans.val() << '\n';
	
	return 0;
}
0