#1116069 (C++23) No.3251 kthmex

提出ソース
結果

問題	No.3251 kthmex
コンテスト
ユーザー	risujiroh
提出日時	2025-08-29 23:02:50
言語	C++23 (gcc 13.3.0 + boost 1.87.0)
結果	RE
実行時間	-
コード長	12,729 bytes
コンパイル時間	6,000 ms
コンパイル使用メモリ	347,424 KB
実行使用メモリ	28,176 KB
最終ジャッジ日時	2025-08-29 23:03:17
合計ジャッジ時間	10,913 ms
ジャッジサーバーID （参考情報）	judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	RE * 1
other	WA * 3 RE * 29
権限があれば一括ダウンロードができます
ソースコード

raw source code
#if __INCLUDE_LEVEL__ == 0

#include __BASE_FILE__

using Mint = atcoder::modint998244353;

void Solve() {
  int n;
  IN(n);
  vector<int> a(n);
  IN(a);
  int A = ranges::max(a);
  HldTree g(n);
  for (int _ : Rep(0, n - 1)) {
    int i, j;
    IN(i, j);
    --i, --j;
    g.add_edge({i, j, 1});
  }
  g.build(0);

  linear_sieve::init(A);
  for (int i : Rep1(2, A)) {
    int prv = -1;
    int len = 0;
    for (int p : linear_sieve::factor(i)) {
      if (p != prv) {
        if (prv != -1) {
          CsrArray<pair<int, int>>::Add(i, {prv, len});
        }
        prv = p;
        len = 1;
      } else {
        ++len;
      }
    }
    CsrArray<pair<int, int>>::Add(i, {prv, len});
  }
  auto factors = CsrArray<pair<int, int>>::Build(A + 1);

  for (int p : linear_sieve::primes) {
    int prod = 1;
    for (int e = 0;; ++e) {
      CsrArray<int>::Add(p, prod);
      if (prod > A / p) {
        break;
      }
      prod *= p;
    }
  }
  auto pw = CsrArray<int>::Build(A + 1);

  vector<Mint> out(n);

  vector<unordered_map<int, int>> f(n);
  vector<Mint> cur(n);
  for (int i : Rev(g.order)) {
    for (auto [p, e] : factors[a[i]]) {
      f[i][p] = e;
    }
    cur[i] = a[i];
    for (auto [j, e] : g.adj[i]) {
      if (g.pe[j] != e) {
        continue;
      }
      if (Sz(f[i]) < Sz(f[j])) {
        swap(f[i], f[j]);
        swap(cur[i], cur[j]);
      }
      for (auto [p, e] : f[j]) {
        auto it = f[i].find(p);
        if (it != f[i].end()) {
          int ei = it->second;
          if (ei < e) {
            it->second = e;
            cur[i] *= pw[p][e - ei];
          } else {
            continue;
          }
        } else {
          f[i][p] = e;
          cur[i] *= pw[p][e];
        }
      }
    }
    out[i] = cur[i];
  }

  ranges::for_each(out, LIFT(OUT));
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  Solve();
}

#elif __INCLUDE_LEVEL__ == 1

#include <bits/stdc++.h>

#include <atcoder/modint.hpp>

template <class T, int Id = -1>
class CsrArray {
 public:
  static void Reserve(int m) {
    buf_.reserve(m);
  }

  static void Add(int i, T x) {
    buf_.emplace_back(i, std::move(x));
  }

  static CsrArray Build(int n) {
    CsrArray ret;
    ret.pos_.resize(n + 1);
    for (int i : buf_ | std::views::keys) {
      ++ret.pos_[i];
    }
    std::partial_sum(ret.pos_.begin(), ret.pos_.end(), ret.pos_.begin());
    ret.data_.resize(ret.pos_[n]);
    for (auto& [i, x] : buf_ | std::views::reverse) {
      ret.data_[--ret.pos_[i]] = std::move(x);
    }
    buf_.clear();
    return ret;
  }

  int size() const { return int(pos_.size()) - 1; }

  auto operator[](int i) {
    return std::span<T>(data_.data() + pos_[i], data_.data() + pos_[i + 1]);
  }
  auto operator[](int i) const {
    return std::span<const T>(data_.data() + pos_[i], data_.data() + pos_[i + 1]);
  }

 private:
  static thread_local inline std::vector<std::pair<int, T>> buf_;
  std::vector<T> data_;
  std::vector<int> pos_;
};

#include <ext/pb_ds/assoc_container.hpp>

struct Splitmix64Hash {
  using u64 = std::uint64_t;

  static u64 splitmix64(u64 x) {
    x += 0x9e3779b97f4a7c15;
    x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9;
    x = (x ^ (x >> 27)) * 0x94d049bb133111eb;
    return x ^ (x >> 31);
  }

  u64 operator()(u64 x) const {
    static const u64 r = std::chrono::steady_clock::now().time_since_epoch().count();
    return splitmix64(x + r);
  }
};

template <class Key, class T>
using HashMap = __gnu_pbds::gp_hash_table<Key, T, Splitmix64Hash>;

template <class Key>
using HashSet = HashMap<Key, __gnu_pbds::null_type>;

template <class T> concept MyRange = std::ranges::range<T> && !std::convertible_to<T, std::string_view>;
template <class T> concept MyTuple = std::__is_tuple_like<T>::value && !MyRange<T>;

namespace std {

istream& operator>>(istream& is, MyRange auto&& r) {
  for (auto&& e : r) is >> e;
  return is;
}
istream& operator>>(istream& is, MyTuple auto&& t) {
  apply([&](auto&... xs) { (is >> ... >> xs); }, t);
  return is;
}

ostream& operator<<(ostream& os, MyRange auto&& r) {
  auto sep = "";
  for (auto&& e : r) os << exchange(sep, " ") << e;
  return os;
}
ostream& operator<<(ostream& os, MyTuple auto&& t) {
  auto sep = "";
  apply([&](auto&... xs) { ((os << exchange(sep, " ") << xs), ...); }, t);
  return os;
}

template <class T, atcoder::internal::is_modint_t<T>* = nullptr>
istream& operator>>(istream& is, T& x) {
  int v;
  is >> v;
  x = T::raw(v);
  return is;
}

template <class T, atcoder::internal::is_modint_t<T>* = nullptr>
ostream& operator<<(ostream& os, const T& x) {
  return os << x.val();
}

}  // namespace std

struct Graph {
  struct Edge {
    int src, dst;
    int64_t cost;

    int other(int v) const {
      __glibcxx_assert(v == src or v == dst);
      return src ^ dst ^ v;
    }
  };

  std::vector<Edge> edges;
  std::vector<std::vector<std::pair<int, int>>> adj;

  Graph() {}
  explicit Graph(int n) : adj(n) {}

  int n() const { return std::size(adj); }
  int m() const { return std::size(edges); }
  int add_edge(const Edge& e, bool directed) {
    __glibcxx_assert(0 <= e.src and e.src < n());
    __glibcxx_assert(0 <= e.dst and e.dst < n());
    int id = m();
    edges.push_back(e);
    adj[e.src].emplace_back(e.dst, id);
    if (not directed) adj[e.dst].emplace_back(e.src, id);
    return id;
  }
};

struct DfsTree : Graph {
  using T = decltype(Edge::cost);

  std::vector<int> root;
  std::vector<int> pv;
  std::vector<int> pe;
  std::vector<int> order;
  std::vector<int> in;
  std::vector<int> out;
  std::vector<int> sub;
  std::vector<int> depth;
  std::vector<int> min_depth;
  std::vector<T> dist;
  std::vector<int> last;
  int num_trials;

  DfsTree() {}
  explicit DfsTree(int n)
      : Graph(n),
        root(n, -1),
        pv(n, -1),
        pe(n, -1),
        in(n, -1),
        out(n, -1),
        sub(n, -1),
        depth(n, -1),
        min_depth(n, -1),
        dist(n, std::numeric_limits<T>::max()),
        last(n, -1),
        num_trials(0) {}

  int add_edge(const Edge& e) { return Graph::add_edge(e, false); }
  void dfs(int r, bool clear_order = true) {
    __glibcxx_assert(0 <= r and r < n());
    root[r] = r;
    pv[r] = -1;
    pe[r] = -1;
    if (clear_order) order.clear();
    depth[r] = 0;
    dist[r] = T{};
    dfs_impl(r);
    ++num_trials;
  }
  void dfs_all() {
    std::fill(std::begin(root), std::end(root), -1);
    for (int v = 0; v < n(); ++v)
      if (root[v] == -1) dfs(v, v == 0);
  }

  int deeper(int id) const {
    __glibcxx_assert(0 <= id and id < m());
    int a = edges[id].src;
    int b = edges[id].dst;
    return depth[a] < depth[b] ? b : a;
  }
  bool is_tree_edge(int id) const {
    __glibcxx_assert(0 <= id and id < m());
    return id == pe[deeper(id)];
  }
  bool is_ancestor(int u, int v) const {
    __glibcxx_assert(0 <= u and u < n());
    __glibcxx_assert(0 <= v and v < n());
    return in[u] <= in[v] and out[v] <= out[u];
  }

 private:
  void dfs_impl(int v) {
    in[v] = std::size(order);
    order.push_back(v);
    sub[v] = 1;
    min_depth[v] = depth[v];
    last[v] = num_trials;
    for (auto&& [u, id] : adj[v]) {
      if (id == pe[v]) continue;
      if (last[u] == num_trials) {
        min_depth[v] = std::min(min_depth[v], depth[u]);
        continue;
      }
      root[u] = root[v];
      pv[u] = v;
      pe[u] = id;
      depth[u] = depth[v] + 1;
      dist[u] = dist[v] + edges[id].cost;
      dfs_impl(u);
      sub[v] += sub[u];
      min_depth[v] = std::min(min_depth[v], min_depth[u]);
    }
    out[v] = std::size(order);
  }
};

struct HldTree : DfsTree {
  std::vector<int> head;

  HldTree() {}
  explicit HldTree(int n) : DfsTree(n), head(n, -1) {}

  void build(int r, bool clear_order = true) {
    __glibcxx_assert(0 <= r and r < n());
    dfs(r, clear_order);
    order.erase(std::end(order) - sub[r], std::end(order));
    head[r] = r;
    build_impl(r);
  }
  void build_all() {
    std::fill(std::begin(root), std::end(root), -1);
    for (int v = 0; v < n(); ++v)
      if (root[v] == -1) build(v, v == 0);
  }

  int lca(int u, int v) const {
    __glibcxx_assert(0 <= u and u < n());
    __glibcxx_assert(0 <= v and v < n());
    __glibcxx_assert(root[u] == root[v]);
    while (true) {
      if (in[u] > in[v]) std::swap(u, v);
      if (head[u] == head[v]) return u;
      v = pv[head[v]];
    }
  }
  int d(int u, int v) const {
    __glibcxx_assert(0 <= u and u < n());
    __glibcxx_assert(0 <= v and v < n());
    __glibcxx_assert(root[u] == root[v]);
    return depth[u] + depth[v] - 2 * depth[lca(u, v)];
  }
  T distance(int u, int v) const {
    __glibcxx_assert(0 <= u and u < n());
    __glibcxx_assert(0 <= v and v < n());
    __glibcxx_assert(root[u] == root[v]);
    return dist[u] + dist[v] - 2 * dist[lca(u, v)];
  }
  int la(int v, int d) const {
    __glibcxx_assert(0 <= v and v < n());
    __glibcxx_assert(0 <= d and d <= depth[v]);
    while (depth[head[v]] > d) v = pv[head[v]];
    return order[in[head[v]] + (d - depth[head[v]])];
  }
  int next(int src, int dst) const {
    __glibcxx_assert(0 <= src and src < n());
    __glibcxx_assert(0 <= dst and dst < n());
    __glibcxx_assert(root[src] == root[dst]);
    __glibcxx_assert(src != dst);
    if (not is_ancestor(src, dst)) return pv[src];
    return la(dst, depth[src] + 1);
  }
  int next(int src, int dst, int k) const {
    __glibcxx_assert(0 <= src and src < n());
    __glibcxx_assert(0 <= dst and dst < n());
    __glibcxx_assert(root[src] == root[dst]);
    __glibcxx_assert(k >= 0);
    int v = lca(src, dst);
    if (k <= depth[src] - depth[v]) return la(src, depth[src] - k);
    k -= depth[src] - depth[v];
    __glibcxx_assert(k <= depth[dst] - depth[v]);
    return la(dst, depth[v] + k);
  }
  template <class Function> void apply(int src, int dst, bool vertex, Function f) const {
    __glibcxx_assert(0 <= src and src < n());
    __glibcxx_assert(0 <= dst and dst < n());
    __glibcxx_assert(root[src] == root[dst]);
    int v = lca(src, dst);
    while (head[src] != head[v]) {
      f(in[src] + 1, in[head[src]]);
      src = pv[head[src]];
    }
    if (vertex)
      f(in[src] + 1, in[v]);
    else if (src != v)
      f(in[src] + 1, in[v] + 1);
    auto rec = [&](auto self, int to) -> void {
      if (head[v] == head[to]) {
        if (v != to) f(in[v] + 1, in[to] + 1);
        return;
      }
      self(self, pv[head[to]]);
      f(in[head[to]], in[to] + 1);
    };
    rec(rec, dst);
  }
  template <class Searcher> int search(int src, int dst, bool vertex, Searcher f) const {
    __glibcxx_assert(0 <= src and src < n());
    __glibcxx_assert(0 <= dst and dst < n());
    __glibcxx_assert(root[src] == root[dst]);
    int res = -1;
    apply(src, dst, vertex, [&](int l, int r) {
      if (res != -1) return;
      int i = f(l, r);
      if (l > r) std::swap(l, r);
      if (l <= i and i < r) res = vertex ? order[i] : pe[order[i]];
    });
    return res;
  }

 private:
  void build_impl(int v) {
    in[v] = std::size(order);
    order.push_back(v);
    auto pos = std::partition(std::begin(adj[v]), std::end(adj[v]), [&](auto&& e) { return e.second == pe[e.first]; });
    auto it =
        std::max_element(std::begin(adj[v]), pos, [&](auto&& a, auto&& b) { return sub[a.first] < sub[b.first]; });
    if (it != std::begin(adj[v])) std::iter_swap(std::begin(adj[v]), it);
    std::partition(pos, std::end(adj[v]), [&](auto&& e) { return e.second == pe[v]; });
    for (auto&& [u, id] : adj[v]) {
      if (id != pe[u]) break;
      head[u] = u == adj[v].front().first ? head[v] : u;
      build_impl(u);
    }
    out[v] = std::size(order);
  }
};

namespace linear_sieve {

std::vector<int> primes, lpf;
void init(int n) {
  if (n < int(std::size(lpf))) return;
  if (n < 2 * int(std::size(lpf))) n = 2 * std::size(lpf);
  lpf.resize(n + 1, -1);
  for (int d = 2; d <= n; ++d) {
    if (lpf[d] == -1) lpf[d] = d, primes.push_back(d);
    for (int p : primes) {
      if (p * d > n or p > lpf[d]) break;
      lpf[p * d] = p;
    }
  }
}
std::vector<int> factor(int n) {
  __glibcxx_assert(n >= 1);
  std::vector<int> res;
  for (init(n); n > 1; n /= res.back()) res.push_back(lpf[n]);
  return res;
}

}  // namespace linear_sieve

using namespace std;

#define _ _ [[maybe_unused]]
#define LIFT(f) ([&](auto&&... xs) -> decltype(auto) { return f(forward<decltype(xs)>(xs)...); })
#define Rev views::reverse
#define Rep(...) [](int l, int r) { return views::iota(min(l, r), r); }(__VA_ARGS__)
#define Rep1(...) [](int l, int r) { return Rep(l, r + 1); }(__VA_ARGS__)
#define Sz(r) int(size(r))
#define IN(...) (cin >> forward_as_tuple(__VA_ARGS__))
#define OUT(...) (cout << forward_as_tuple(__VA_ARGS__) << '\n')

#endif  // __INCLUDE_LEVEL__ == 1
yukicoder

結果

ソースコード
