CompetitiveProgrammingCpp

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub

:heavy_check_mark: Test/Graph/Tree/ReRootingDP.test.cpp

Depends on

Code

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

#include <iostream>

// begin:tag includes
#include "../../../Library/Graph/Tree/ReRootingDP.hpp"
#include "../../../Library/Math/ModInt.hpp"
// end:tag includes

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.addEdge(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;
}
#line 1 "Test/Graph/Tree/ReRootingDP.test.cpp"
#include <vector>
#define PROBLEM "https://yukicoder.me/problems/no/2360"

#include <iostream>

// begin:tag includes
#line 2 "Library/Graph/Tree/ReRootingDP.hpp"
#include <unordered_map>

#line 4 "Library/Graph/Tree/ReRootingDP.hpp"

#line 2 "Library/Algebraic/Monoid.hpp"

#line 4 "Library/Algebraic/Monoid.hpp"

namespace mtd {

  template <class S,    // set
            S element,  // identity element
            class op    // binary operation
            >
  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_monoid_specialization_of = requires {
      typename std::enable_if_t<std::is_same_v<
          T, S<typename T::value_type, T::_element, typename T::op_type>>>;
    };
  }  // namespace __detail

  template <typename M>
  concept monoid = __detail::is_monoid_specialization_of<M, Monoid>;

}  // namespace mtd
#line 2 "Library/Graph/Normal/BFS.hpp"

#include <concepts>

#include <queue>

#line 6 "Library/Graph/Normal/BFS.hpp"

#line 2 "Library/Graph/Graph.hpp"
#include <deque>

#line 4 "Library/Graph/Graph.hpp"
#include <ranges>

#include <tuple>

#line 7 "Library/Graph/Graph.hpp"

namespace mtd {
  template <class Node = long long, 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) {}
    Graph(int n, const std::vector<std::tuple<Node, Node>>& edges,
          bool is_arc = false, bool is_index1 = true)
        : Graph<Node, Cost>(n) {
      for (auto [u, v] : edges) {
        u -= is_index1;
        v -= is_index1;
        if (is_arc) {
          addArc(u, v);
        } else {
          addEdge(u, v);
        }
      }
    }
    Graph(int n, const std::vector<std::tuple<Node, Node, Cost>>& edges,
          bool is_arc = false, bool is_index1 = true)
        : Graph<Node, Cost>(n) {
      for (auto [u, v, c] : edges) {
        u -= is_index1;
        v -= is_index1;
        if (is_arc) {
          addArc(u, v, c);
        } else {
          addEdge(u, v, c);
        }
      }
    }

    auto addEdge(const Node& f, const Node& t, const Cost& c = 1) {
      addArc(f, t, c);
      addArc(t, f, c);
    }
    auto addArc(const Node& f, const Node& t, const Cost& c = 1) {
      m_graph[f].emplace_back(t, 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.addArc(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

#line 8 "Library/Graph/Normal/BFS.hpp"

namespace mtd {
  template <class Node, class Cost, class Lambda,
            std::convertible_to<Node> _Node>
  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

#line 2 "Library/Graph/Tree/TreeDP.hpp"

#line 6 "Library/Graph/Tree/TreeDP.hpp"

#line 8 "Library/Graph/Tree/TreeDP.hpp"

namespace mtd {
  template <class Node, class Cost, class Lambda,
            std::convertible_to<Node> _Node>
  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

#line 8 "Library/Graph/Tree/ReRootingDP.hpp"

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();

    // <辺情報を考慮したMonoidの2項演算>

    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));
    };

    // <node:toを根とした木で全てマージした解を求める>

    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);
    };

    // <node:toを根とした木でfrom以外マージした解を求める>

    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);
      }
      // node情報を反映させて値を確定

      ret_m[to] = node_f(val_ord, to);
      for (auto&& [_, pac] : partial_ac[to]) { pac = node_f(pac, to); }
    };

    // rootを根とした解を求める

    treeDP(graph, root, [&](Node f, Node t, const Cost& c) {
      partial[t].emplace_back(all_merge(f), f, c);
    });
    accumulation(0);

    // rootからbfsして各nodeを根とした解を求める

    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;
  }
}  // namespace mtd

#line 2 "Library/Math/ModInt.hpp"

#line 4 "Library/Math/ModInt.hpp"
#include <iterator>

#line 2 "Library/Math/Math.hpp"

#include <cmath>

#include <numeric>

#include <optional>

#line 9 "Library/Math/Math.hpp"

#line 2 "Library/Math/EuclideanAlgorithm.hpp"

#line 6 "Library/Math/EuclideanAlgorithm.hpp"

namespace mtd {

  class EuclideanAlgorithm {
    using T = long long;

    // 大きすぎるとオーバーフローしてしまう

    const static inline T m_mx = 1e9;

    const T m_a;
    const T m_b;
    const T m_c;

    T m_gcd;
    T m_x;
    T m_y;

    auto excludedEuclidAlgorithm(T a, T b) -> std::tuple<T, T, T> {
      if (a < 0) {
        auto [g, x, y] = excludedEuclidAlgorithm(-a, -b);
        return {g, -x, -y};
      }
      if (b == 0) { return {a, 1, 0}; }
      auto [g, y, x] = excludedEuclidAlgorithm(b, a % b);
      y -= a / b * x;
      return {g, x, y};
    }

    auto kRange(T x, T b, T l) const -> std::pair<T, T> {
      // x + b * k >= l を満たす k の範囲を求める

      T xd = (l - x);
      if (b == 0 && x >= l) { return {-m_mx, m_mx}; }
      if (b == 0 && x < l) { return {m_mx, -m_mx}; }
      if (b > 0 && xd < 0) { return {xd / b, m_mx}; }
      if (b > 0 && xd >= 0) { return {(xd + b - 1) / b, m_mx}; }
      if (b < 0 && xd < 0) { return {-m_mx, (-xd) / (-b)}; }
      if (b < 0 && xd >= 0) { return {-m_mx, -(xd - b - 1) / (-b)}; }
      return {m_mx, -m_mx};
    }

  public:
    auto debug() const {
      std::cout << m_a << " * " << m_x << " + " << m_b << " * " << m_y << " = "
                << m_c << std::endl;
      std::cout << "calc: " << m_a * m_x + m_b * m_y << " = " << m_c
                << std::endl;
    }

    EuclideanAlgorithm(T a, T b, T c) : m_a(a), m_b(b), m_c(c) {
      if (a == 0 && b == 0) { throw std::runtime_error(""); }
      auto [g, x, y] = excludedEuclidAlgorithm(a, b);
      if (c % g > 0) {
        throw std::runtime_error(
            "There is no solution to the equation. c must be divisible by "
            "gcd(a,b).");
      }
      m_gcd = g;
      m_x = c / g * x;
      m_y = c / g * y;
    }
    EuclideanAlgorithm(T a, T b) : EuclideanAlgorithm(a, b, std::gcd(a, b)) {}

    auto gcd() const { return m_gcd; }
    auto get(T x, T y) const { return m_a * x + m_b * y; }
    auto get(T k) const -> std::pair<T, T> {
      if (m_b == 0) { return {m_x, m_y - k}; }
      if (m_a == 0) { return {m_x + k, m_y}; }
      return {m_x + m_b * k, m_y - m_a * k};
    }
    // x>=x_lとなるようなkの範囲

    auto getMinX(T x_l = 0) const -> std::pair<T, T> {
      return kRange(m_x, m_b, x_l);
    }
    // y>=y_lとなるようなkの範囲

    auto getMinY(T y_l = 0) const -> std::pair<T, T> {
      return kRange(m_y, -1 * m_a, y_l);
    }
    // x>=x_l, y>=y_lとなるようなkの範囲

    auto getMin(T x_l = 0, T y_l = 0) const -> std::pair<T, T> {
      auto [xl, xr] = getMinX(x_l);
      auto [yl, yr] = getMinY(y_l);
      return {std::max(xl, yl), std::min(xr, yr)};
    }
  };

}  // namespace mtd

#line 11 "Library/Math/Math.hpp"

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(int size = 3 * static_cast<int>(1e6))
        : m_fac(constructFac(size)), m_finv(constructInv(size)) {}

    /* O(log b) */
    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;
    }

    /* O(log mod) */
    template <class S>
    static constexpr std::optional<long long> log(S x, S y, S mod) {
      x %= mod;
      y %= mod;

      if (mod == 1) { return 0; }
      if (x == 0 && y == 0) { return 1; }
      if (x == 0 && y == 1) { return 0; }
      if (x == 0) { return std::nullopt; }
      if (y == 1) { return 0; }

      if (auto g = std::gcd(x, mod); g > 1) {
        if (y % g) { return std::nullopt; }
        auto nx = x / g;
        auto nmod = mod / g;
        auto ea = mtd::EuclideanAlgorithm(nx, -nmod, 1);
        auto [t, _] = ea.getMinX();
        auto [nx_inv, __] = ea.get(t);
        nx_inv %= nmod;
        if (auto ans = log(x, y / g * nx_inv, nmod); ans) {
          return ans.value() + 1;
        } else {
          return ans;
        }
      }

      auto s = static_cast<S>(std::sqrt(mod));
      S xe = y;
      std::unordered_map<S, S> map;
      map.reserve(s);
      for (auto i : std::views::iota(0, s)) {
        (xe *= x) %= mod;
        map[xe] = i + 1;
      }

      S xs = 1;
      for ([[maybe_unused]] auto _ : std::views::iota(0, s)) {
        (xs *= x) %= mod;
      }
      S xse = 1;
      for (auto i : std::views::iota(0, mod / s + 5)) {
        (xse *= xs) %= mod;
        if (map.contains(xse)) { return s * (i + 1) - map[xse]; }
      }
      return std::nullopt;
    }
    constexpr std::optional<long long> log(long long x,
                                           long long y) requires requires {
      typename T::value_type;
      T::mod();
    }
    { return log(x, y, T::mod()); }

    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

#line 7 "Library/Math/ModInt.hpp"

namespace mtd {

  template <int MOD, class T = long long>
  class ModInt {
  public:
    using value_type = T;
    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; }
    static constexpr auto mod() { return MOD; }
  };

}  // namespace mtd
#line 9 "Test/Graph/Tree/ReRootingDP.test.cpp"
// end:tag includes

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.addEdge(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;
}
Back to top page