1254D - Tree Queries - CodeForces Solution

#include <algorithm>
#include <cassert>
#include <cmath>
#include <iostream>
#include <vector>
using namespace std;
const int MOD = 998244353;
struct mod_int {
  int val;
  mod_int(long long v = 0) {
    if (v < 0)
      v = v % MOD + MOD;
    if (v >= MOD)
      v %= MOD;
    val = v;
  }
  static int mod_inv(int a, int m = MOD) {
    int g = m, r = a, x = 0, y = 1;
    while (r != 0) {
      int q = g / r;
      g %= r;
      swap(g, r);
      x -= q * y;
      swap(x, y);
    }
    return x < 0 ? x + m : x;
  }
  explicit operator int() const { return val; }
  mod_int &operator+=(const mod_int &other) {
    if ((val += other.val) >= MOD)
      val -= MOD;
    return *this;
  }
  mod_int &operator-=(const mod_int &other) {
    if ((val -= other.val) < 0)
      val += MOD;
    return *this;
  }
  static unsigned int fast_mod(uint64_t x, unsigned int m = MOD) {
    return x % m;
  }
  mod_int &operator*=(const mod_int &other) {
    val = fast_mod((uint64_t)val * other.val);
    return *this;
  }
  mod_int &operator/=(const mod_int &other) { return *this *= other.inv(); }
  friend mod_int operator+(const mod_int &a, const mod_int &b) {
    return mod_int(a) += b;
  }
  friend mod_int operator-(const mod_int &a, const mod_int &b) {
    return mod_int(a) -= b;
  }
  friend mod_int operator*(const mod_int &a, const mod_int &b) {
    return mod_int(a) *= b;
  }
  friend mod_int operator/(const mod_int &a, const mod_int &b) {
    return mod_int(a) /= b;
  }
  mod_int &operator++() {
    val = val == MOD - 1 ? 0 : val + 1;
    return *this;
  }
  mod_int &operator--() {
    val = val == 0 ? MOD - 1 : val - 1;
    return *this;
  }
  mod_int operator++(int) {
    mod_int before = *this;
    ++*this;
    return before;
  }
  mod_int operator--(int) {
    mod_int before = *this;
    --*this;
    return before;
  }
  mod_int operator-() const { return val == 0 ? 0 : MOD - val; }
  bool operator==(const mod_int &other) const { return val == other.val; }
  bool operator!=(const mod_int &other) const { return val != other.val; }
  mod_int inv() const { return mod_inv(val); }
  mod_int pow(long long p) const {
    assert(p >= 0);
    mod_int a = *this, result = 1;
    while (p > 0) {
      if (p & 1)
        result *= a;
      a *= a;
      p >>= 1;
    }
    return result;
  }
  friend ostream &operator<<(ostream &stream, const mod_int &m) {
    return stream << m.val;
  }
};
template <typename T, bool maximum_mode = false> struct RMQ {
  int n = 0, levels = 0;
  vector<T> values;
  vector<vector<int>> range_low;
  RMQ(const vector<T> &_values = {}) {
    if (!_values.empty())
      build(_values);
  }
  static int largest_bit(int x) { return 31 - __builtin_clz(x); }
  int better_index(int a, int b) const {
    return (values[a] < values[b]) ^ maximum_mode ? a : b;
  }
  void build(const vector<T> &_values) {
    values = _values;
    n = values.size();
    levels = largest_bit(n) + 1;
    range_low.resize(levels);
    for (int k = 0; k < levels; k++)
      range_low[k].resize(n - (1 << k) + 1);
    for (int i = 0; i < n; i++)
      range_low[0][i] = i;
    for (int k = 1; k < levels; k++)
      for (int i = 0; i <= n - (1 << k); i++)
        range_low[k][i] = better_index(range_low[k - 1][i],
                                       range_low[k - 1][i + (1 << (k - 1))]);
  }
  int query_index(int a, int b) const {
    assert(0 <= a && a < b && b <= n);
    int level = largest_bit(b - a);
    return better_index(range_low[level][a],
                        range_low[level][b - (1 << level)]);
  }
  T query_value(int a, int b) const { return values[query_index(a, b)]; }
};
struct LCA {
  int n = 0;
  vector<vector<int>> adj;
  vector<int> parent, depth, subtree_size;
  vector<int> euler, first_occurrence;
  vector<int> tour_start, tour_end, tour_list;
  RMQ<int> rmq;
  LCA(int _n = 0) { init(_n); }
  LCA(const vector<vector<int>> &_adj) { init(_adj); }
  void init(int _n) {
    n = _n;
    adj.assign(n, {});
    parent.resize(n);
    depth.resize(n);
    subtree_size.resize(n);
    first_occurrence.resize(n);
    tour_start.resize(n);
    tour_end.resize(n);
    tour_list.resize(n);
  }
  void init(const vector<vector<int>> &_adj) {
    init(_adj.size());
    adj = _adj;
  }
  void add_edge(int a, int b) {
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  int tour;
  void dfs(int node, int par) {
    parent[node] = par;
    depth[node] = par < 0 ? 0 : depth[par] + 1;
    subtree_size[node] = 1;
    first_occurrence[node] = euler.size();
    euler.push_back(node);
    tour_list[tour] = node;
    tour_start[node] = tour++;
    if (par >= 0)
      adj[node].erase(find(adj[node].begin(), adj[node].end(), par));
    for (int child : adj[node]) {
      dfs(child, node);
      subtree_size[node] += subtree_size[child];
      euler.push_back(node);
    }
    tour_end[node] = tour;
  }
  void build() {
    tour = 0;
    parent.assign(n, -1);
    for (int i = 0; i < n; i++)
      if (parent[i] < 0) {
        dfs(i, -1);
        euler.push_back(-1);
      }
    assert((int)euler.size() == 2 * n);
    vector<int> euler_depths;
    for (int node : euler)
      euler_depths.push_back(node < 0 ? node : depth[node]);
    rmq.build(euler_depths);
  }
  int get_lca(int a, int b) const {
    a = first_occurrence[a];
    b = first_occurrence[b];
    if (a > b)
      swap(a, b);
    return euler[rmq.query_index(a, b + 1)];
  }
  bool is_ancestor(int a, int b) const {
    return tour_start[a] <= tour_start[b] && tour_start[b] < tour_end[a];
  }
  bool on_path(int x, int a, int b) const {
    int anc = get_lca(a, b);
    return is_ancestor(anc, x) && (is_ancestor(x, a) || is_ancestor(x, b));
  }
  int get_dist(int a, int b) const {
    return depth[a] + depth[b] - 2 * depth[get_lca(a, b)];
  }
  int child_ancestor(int a, int b) const {
    assert(a != b);
    assert(is_ancestor(a, b));
    int child =
        euler[rmq.query_index(first_occurrence[a], first_occurrence[b] + 1) +
              1];
    assert(is_ancestor(child, b));
    return child;
  }
};
struct query {
  int type, v;
  mod_int d;
};
int N, APPLY, Q;
LCA lca;
vector<mod_int> values, updates, query_d;
int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cin >> N >> Q;
  APPLY = 1.0 * sqrt(N) + 1;
  cerr << "Apply size: " << APPLY << endl;
  lca.init(N);
  for (int i = 0; i < N - 1; i++) {
    int u, v;
    cin >> u >> v;
    u--;
    v--;
    lca.add_edge(u, v);
  }
  lca.build();
  values.assign(N, 0);
  vector<query> pending_queries;
  mod_int inv_N = mod_int(1) / N;
  for (int q = 0; q < Q; q++) {
    query qry;
    cin >> qry.type >> qry.v;
    qry.v--;
    if (qry.type == 1) {
      int d;
      cin >> d;
      qry.d = d;
      pending_queries.push_back(qry);
    } else if (qry.type == 2) {
      mod_int answer = values[qry.v];
      for (query &prev : pending_queries)
        if (prev.type == 1) {
          if (prev.v == qry.v) {
            answer += prev.d;
          } else if (lca.is_ancestor(prev.v, qry.v)) {
            int ca = lca.child_ancestor(prev.v, qry.v);
            mod_int probability = (N - lca.subtree_size[ca]) * inv_N;
            answer += probability * prev.d;
          } else {
            mod_int probability = lca.subtree_size[prev.v] * inv_N;
            answer += probability * prev.d;
          }
        }
      cout << answer << '\n';
    } else {
      assert(false);
    }
    if ((int)pending_queries.size() >= APPLY) {
      updates.assign(N, 0);
      query_d.assign(N, 0);
      for (query &prev : pending_queries)
        query_d[prev.v] += prev.d;
      for (int node = 0; node < N; node++)
        if (query_d[node] != 0) {
          mod_int base_probability = lca.subtree_size[node] * inv_N;
          updates[0] += base_probability * query_d[node];
          updates[node] += (1 - base_probability) * query_d[node];
          for (int child : lca.adj[node])
            updates[child] -= lca.subtree_size[child] * inv_N * query_d[node];
        }
      for (int node : lca.tour_list) {
        values[node] += updates[node];
        for (int child : lca.adj[node])
          updates[child] += updates[node];
      }
      pending_queries.clear();
    }
  }
}
// dQjruxkazRXIsLLKAamQUeAIwBWFDqwPrsaHzPeuHgrMDnsIuWQkyyAFAIASwEnZbYBdiqWIedCxkHsOmSGFDBtFVhjOChWsDSXoDiyeUwDBObMgCLCOeCXsRpFwhKAR