結果

問題 No.2360 Path to Integer
ユーザー cutmdocutmdo
提出日時 2024-12-19 23:41:01
言語 C++23
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 270 ms / 2,500 ms
コード長 11,734 bytes
コンパイル時間 2,604 ms
コンパイル使用メモリ 166,052 KB
実行使用メモリ 102,284 KB
最終ジャッジ日時 2024-12-19 23:41:09
合計ジャッジ時間 6,507 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 45 ms
49,944 KB
testcase_01 AC 45 ms
50,084 KB
testcase_02 AC 44 ms
49,944 KB
testcase_03 AC 45 ms
50,136 KB
testcase_04 AC 45 ms
50,056 KB
testcase_05 AC 45 ms
50,096 KB
testcase_06 AC 45 ms
50,072 KB
testcase_07 AC 46 ms
50,588 KB
testcase_08 AC 63 ms
55,100 KB
testcase_09 AC 270 ms
99,664 KB
testcase_10 AC 250 ms
102,284 KB
testcase_11 AC 247 ms
101,608 KB
testcase_12 AC 186 ms
100,912 KB
testcase_13 AC 266 ms
99,528 KB
testcase_14 AC 244 ms
98,480 KB
testcase_15 AC 260 ms
99,396 KB
testcase_16 AC 245 ms
98,352 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <vector>
#define PROBLEM "https://yukicoder.me/problems/no/2360"

#include <iostream>

#include <queue>
#include <unordered_map>
#include <ranges>
#include <tuple>
#include <vector>
#include <iterator>
#include <iostream>
#include <deque>
namespace mtd {  template <class S,                S element,              class op                >  requires std::is_invocable_r_v<S, op, S, S>  struct Monoid {    using value_type = S;    constexpr static S _element = element;    using op_type = op;    S m_val;    constexpr Monoid(S val) : m_val(val) {}    constexpr Monoid() : Monoid(element) {}    constexpr Monoid binaryOperation(const Monoid& m2) const {      return op()(m_val, m2.m_val);    }    friend std::ostream& operator<<(std::ostream& os,                                    const Monoid<S, element, op>& m) {      return os << m.m_val;    }  };  namespace __detail {    template <typename T, template <typename, auto, typename> typename S>    concept is_specialization_of = requires {      typename std::enable_if_t<std::is_same_v<          T, S<typename T::value_type, T::_element, typename T::op_type>>>;    };  }    template <typename M>  concept monoid = __detail::is_specialization_of<M, Monoid>;}  
namespace mtd {  template <class Node = int, class Cost = long long>  class Graph {            using Edge = std::pair<Node, Cost>;    using Edges = std::vector<Edge>;    const int m_n;    std::vector<Edges> m_graph;  public:    Graph(int n) : m_n(n), m_graph(n) {}    Graph(const std::vector<Edges>& edges)        : m_n(edges.size()), m_graph(edges) {}    auto addEdge(const Node& f, const Node& t, const Cost& c = 1) {      m_graph[f].emplace_back(t, c);    }    auto addEdgeUndirected(const Node& f, const Node& t, const Cost& c = 1) {      addEdge(f, t, c);      addEdge(t, f, c);    }    auto getEdges(const Node& from) const {      class EdgesRange {        const typename Edges::const_iterator b, e;      public:        EdgesRange(const Edges& edges) : b(edges.begin()), e(edges.end()) {}        auto begin() const { return b; }        auto end() const { return e; }      };      return EdgesRange(m_graph[from]);    }    auto getEdges() const {      std::deque<std::tuple<Node, Node, Cost>> edges;      for (Node from : std::views::iota(0, m_n)) {        for (const auto& [to, c] : getEdges(from)) {          edges.emplace_back(from, to, c);        }      }      return edges;    }    auto getEdgesExcludeCost() const {      std::deque<std::pair<Node, Node>> edges;      for (Node from : std::views::iota(0, m_n)) {        for (const auto& [to, _] : getEdges(from)) {          edges.emplace_back(from, to);        }      }      return edges;    }    auto reverse() const {      auto rev = Graph<Node, Cost>(m_n);      for (const auto& [from, to, c] : getEdges()) { rev.addEdge(to, from, c); }      return rev;    }    auto size() const { return m_n; };    auto debug(bool directed = false) const {      for (const auto& [f, t, c] : getEdges()) {        if (f < t || directed) {          std::cout << f << " -> " << t << ": " << c << std::endl;        }      }    }  };}  
namespace mtd {  template <class T>  class Math {    const std::vector<T> m_fac;    const std::vector<T> m_finv;    auto constructFac(int s) {      std::vector<T> fac(s);      fac[0] = fac[1] = 1;      for (long long i = 2; i < s; ++i) { fac[i] = fac[i - 1] * i; }      return fac;    }    auto constructInv(int s) {      std::vector<T> finv(s);      finv[s - 1] = 1 / m_fac[s - 1];      for (long long i = s - 2; i >= 0; --i) {        finv[i] = finv[i + 1] * (i + 1);      }      return finv;    }  public:    constexpr Math(long long size = 3 * 1e6)        : m_fac(constructFac(size)), m_finv(constructInv(size)) {}    static constexpr T pow(T a, long long b) {      T ans = 1;      while (b > 0) {        if (b & 1) { ans *= a; }        b >>= 1;        a *= a;      }      return ans;    }    constexpr auto fact(int n) const { return (n < 0) ? 0 : m_fac[n]; }    constexpr auto factInv(int n) const { return (n < 0 ? 0 : m_finv[n]); }    constexpr auto comb(int n, int r) const {      return fact(n) * factInv(r) * factInv(n - r);    }    constexpr auto perm(int n, int r) const { return fact(n) * factInv(n - r); }  };}  
namespace mtd {  template <class Node, class Cost, class Lambda>  auto bfs(const Graph<Node, Cost>& graph, const Node& root,           const Lambda& lambda) {    auto n = graph.size();    std::vector<bool> used(n);    used[root] = true;    std::queue<Node> q;    q.emplace(root);    while (!q.empty()) {      auto from = q.front();      q.pop();      for (const auto& [to, cost] : graph.getEdges(from)) {        if (used[to]) { continue; }        q.emplace(to);        used[to] = true;        lambda(from, to, cost);      }    }  }}  
namespace mtd {  template <class Node, class Cost, class Lambda>  auto treeDP(const Graph<Node, Cost>& tree, Node root, const Lambda& lambda) {    auto n = tree.size();    std::vector<Node> in(n);    for (const auto& [f, t] : tree.getEdgesExcludeCost())      if (f < t) {        ++in[f];        ++in[t];      }    std::queue<Node> q;    std::vector<bool> used(n);    for (Node i = 0; i < n; ++i)      if (i != root && in[i] == 1) { q.emplace(i); }    while (!q.empty()) {      auto from = q.front();      q.pop();      used[from] = true;      for (const auto& [to, cost] : tree.getEdges(from)) {        if (used[to]) { continue; }        lambda(from, to, cost);        --in[to];        if (to != root && in[to] == 1) { q.emplace(to); }      }    }  }}  
namespace mtd {  template <int MOD, class T = long long>  class ModInt {  public:    T x;    constexpr ModInt(T x) : x(x >= 0 ? x % MOD : MOD + (x % MOD)) {}    constexpr ModInt() : ModInt(0) {}        constexpr auto& operator+=(const ModInt<MOD, T>& m) {      x += m.x;      if (x >= MOD) { x -= MOD; }      return *this;    }    constexpr auto& operator-=(const ModInt<MOD, T>& m) {      x -= m.x;      if (x < 0) { x += MOD; }      return *this;    }    constexpr auto& operator*=(const ModInt<MOD, T>& m) {      x *= m.x;      if (x >= MOD) { x %= MOD; }      return *this;    }    constexpr auto& operator/=(const ModInt<MOD, T>& m) {      x *= mtd::Math<ModInt<MOD, T>>::pow(m.x, MOD - 2).x;      if (x >= MOD) { x %= MOD; }      return *this;    }    constexpr auto operator+(const ModInt<MOD, T>& m) const {      auto t = *this;      t += m;      return t;    }    constexpr auto operator-(const ModInt<MOD, T>& m) const {      auto t = *this;      t -= m;      return t;    }    constexpr auto operator*(const ModInt<MOD, T>& m) const {      auto t = *this;      t *= m;      return t;    }    constexpr auto operator/(const ModInt<MOD, T>& m) const {      auto t = *this;      t /= m;      return t;    }    constexpr auto& operator+=(const T& t) {      return *this += ModInt<MOD, T>(t);    }    constexpr auto& operator-=(const T& t) {      return *this -= ModInt<MOD, T>(t);    }    constexpr auto& operator*=(const T& n) {      return *this *= ModInt<MOD, T>(n);    }    constexpr auto& operator/=(const T& n) {      return *this /= ModInt<MOD, T>(n);    }    constexpr auto operator+(const T& t) const {      return *this + ModInt<MOD, T>(t);    }    constexpr auto operator-(const T& t) const {      return *this - ModInt<MOD, T>(t);    }    constexpr auto operator*(const T& t) const {      return *this * ModInt<MOD, T>(t);    }    constexpr auto operator/(const T& t) const {      return *this / ModInt<MOD, T>(t);    }    constexpr friend auto operator+(const T& t, const ModInt<MOD, T>& m) {      return m + t;    }    constexpr friend auto operator-(const T& t, const ModInt<MOD, T>& m) {      return -m + t;    }    constexpr friend auto operator*(const T& t, const ModInt<MOD, T>& m) {      return m * t;    }    constexpr friend auto operator/(const T& t, const ModInt<MOD, T>& m) {      return ModInt<MOD, T>(1) / m * t;    }        constexpr auto operator-() const { return ModInt<MOD, T>(0 - x); }        constexpr auto operator!=(const ModInt<MOD, T>& m) const {      return x != m.x;    }    constexpr auto operator==(const ModInt<MOD, T>& m) const {      return !(x != m.x);    }        constexpr friend std::ostream& operator<<(std::ostream& os,                                              const ModInt<MOD, T>& m) {      return os << m.x;    }    constexpr friend std::istream& operator>>(std::istream& is,                                              ModInt<MOD, T>& m) {      return is >> m.x;    }    constexpr auto val() const { return x; }  };}  
namespace mtd {  /*   * Monoid: 部分木の情報   * edge_f: 辺の情報を親に流す関数: (M, f, t, c) -> M   * node_f: 子の情報を親に流す関数: (M, i) -> M   */  template <monoid Monoid, class Node, class Cost, class Lambda1, class Lambda2>  auto reRootingDP(const Graph<Node, Cost>& graph, const Lambda1& edge_f,                   const Lambda2& node_f) {    constexpr int root = 0;    auto n = graph.size();        auto merge = [&](Monoid& m1, const Monoid& m2, Node f = -1, Node t = -1,                     const Cost& c = Cost()) {      m1 = m1.binaryOperation((f == -1 || t == -1) ? m2 : edge_f(m2, f, t, c));    };        std::vector<std::vector<std::tuple<Monoid, Node, Cost>>> partial(n);    auto all_merge = [&](Node to) {      Monoid val{};      for (const auto& [ad, from, cost] : partial[to]) {        merge(val, ad, from, to, cost);      }      return node_f(val, to);    };        std::vector<std::unordered_map<Node, Monoid>> partial_ac(n);    std::vector<Monoid> ret_m(n);    auto accumulation = [&](Node to) {            Monoid val_ord{};      for (const auto& [ad, from, cost] : partial[to]) {        partial_ac[to].emplace(from, val_ord);        merge(val_ord, ad, from, to, cost);      }            Monoid val_rord{};      for (auto rit = partial[to].rbegin(); rit != partial[to].rend(); ++rit) {        auto [ad, from, cost] = *rit;        merge(partial_ac[to][from], val_rord, cost);        merge(val_rord, ad, from, to, cost);      }            ret_m[to] = node_f(val_ord, to);      for (auto&& [_, pac] : partial_ac[to]) { pac = node_f(pac, to); }    };        treeDP(graph, root, [&](Node f, Node t, const Cost& c) {      partial[t].emplace_back(all_merge(f), f, c);    });    accumulation(0);        bfs(graph, root, [&](Node f, Node t, const Cost& c) {      partial[t].emplace_back(partial_ac[f][t], f, c);      accumulation(t);    });    std::vector<typename Monoid::value_type> ret;    for (const auto x : ret_m) { ret.emplace_back(x.m_val); }    return ret;  }}  

using ll = long long;
constexpr ll MOD = 998244353;
using mint = mtd::ModInt<MOD>;
auto math = mtd::Math<mint>();

struct S {
  mint m;
  ll s;
  constexpr S(const mint& m, const ll s) : m(m), s(s) {}
  constexpr S() : S(0, 0) {}
};

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(0);

  ll n;
  std::cin >> n;
  std::vector<ll> a(n);
  for (auto i : std::views::iota(0, n)) { std::cin >> a[i]; }
  auto graph = mtd::Graph<>(n);
  for ([[maybe_unused]] auto _ : std::views::iota(0, n - 1)) {
    int u, v;
    std::cin >> u >> v;
    graph.addEdgeUndirected(u - 1, v - 1);
  }

  std::vector<mint> at;
  for (auto x : a) {
    auto size = std::to_string(x).size();
    at.emplace_back(math.pow(10, size));
  }

  auto op = [](const S& a, const S& b) { return S(a.m + b.m, a.s + b.s); };
  using M = mtd::Monoid<S, S(0, 0), decltype(op)>;

  auto edge_f = [&](const M& m, int f, int t, int c) {
    return M(S(m.m_val.m * at[t] + m.m_val.s * a[t], m.m_val.s));
  };
  auto node_f = [&](const M& m, int i) {
    return M(S(m.m_val.m + a[i], m.m_val.s + 1));
  };
  auto dp = mtd::reRootingDP<M>(graph, edge_f, node_f);

  mint ans = 0;
  for (auto x : dp) { ans += x.m; }
  std::cout << ans << std::endl;
}
0