#726489 (C++17) No.1796 木上のクーロン

提出ソース
結果

問題	No.1796 木上のクーロン
コンテスト
ユーザー	Kude
提出日時	2021-12-25 17:35:18
言語	C++17 (gcc 13.3.0 + boost 1.87.0)
結果	AC
実行時間	1,787 ms / 10,000 ms
コード長	6,560 bytes
コンパイル時間	3,147 ms
コンパイル使用メモリ	236,152 KB
最終ジャッジ日時	2025-01-27 06:34:38
ジャッジサーバーID （参考情報）	judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 3
other	AC * 34
権限があれば一括ダウンロードができます
ソースコード

raw source code
#include<bits/stdc++.h>
namespace {
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-function"
#include<atcoder/all>
#pragma GCC diagnostic pop
using namespace std;
using namespace atcoder;
#define rep(i,n)for (int i = 0; i < int(n); ++i)
#define rrep(i,n)for (int i = int(n)-1; i >= 0; --i)
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
template<class T> void chmax(T& a, const T& b) { a = max(a, b); }
template<class T> void chmin(T& a, const T& b) { a = min(a, b); }
using ll = long long;
using P = pair<int,int>;
using VI = vector<int>;
using VVI = vector<VI>;
using VL = vector<ll>;
using VVL = vector<VL>;
using mint = modint998244353;

template<int element_size>
struct array_set {
  int state[element_size];  // -1 if not in set else index in elems
  int elems[element_size];
  int size = 0;
  array_set() { memset(state, -1, sizeof(int) * element_size); }
  void insert(int x) {
    if (state[x] == -1) {
      state[x] = size;
      elems[size] = x;
      size++;
    }
  }
  bool contains(int x) { return state[x] != -1; }
  int* begin() { return elems; }
  int* end() { return elems + size; }
  void clear() { while(size) state[elems[--size]] = -1; }
  void erase(int x) {  // not tested
    int idx = state[x];
    if (idx != -1) {
      state[x] = -1;
      size--;
      if (idx != size) {
        int y = elems[size];
        state[y] = idx;
        elems[idx] = y;
      }
    }
  }
};
template<class T, int element_size>
struct array_unordered_map {
  int state[element_size];  // -1 if not in set else index in elems
  pair<int, T> elems[element_size];
  int size = 0;
  array_unordered_map() { memset(state, -1, sizeof(int) * element_size); }
  bool contains(int x) { return state[x] != -1; }
  pair<int, T>* begin() { return elems; }
  pair<int, T>* end() { return elems + size; }
  void clear() { while(size) state[elems[--size].first] = -1; }
  void erase(int x) {  // not tested
    int idx = state[x];
    if (idx != -1) {
      state[x] = -1;
      size--;
      if (idx != size) {
        int y = elems[size].first;
        state[y] = idx;
        elems[idx] = move(elems[size]);
      }
    }
  }
  T& operator[](int x) {
    if (state[x] == -1) {
      state[x] = size;
      elems[size].first = x;
      elems[size].second = T();
      size++;
    }
    return elems[state[x]].second;
  }
};

constexpr int FACT_SIZE = 1000000;
mint Fact[FACT_SIZE + 1];
mint iFact[FACT_SIZE + 1];
const auto fact_init = [] {
    Fact[0] = mint::raw(1);
    for(int i = 1; i <= FACT_SIZE; ++i) {
        Fact[i] = Fact[i-1] * i;
    }
    iFact[FACT_SIZE] = Fact[FACT_SIZE].inv();
    for(int i = FACT_SIZE; i; --i) {
        iFact[i-1] = iFact[i] * i;
    }
    return false;
}();

mint comb(int n, int k) {
    if (k == 0) return mint::raw(1);
    assert(n >= 0 && k >= 0);
    if (k > n) return mint::raw(0);
    return Fact[n] * iFact[n - k] * iFact[k];
}

mint icomb(int n, int k) {
    return iFact[n] * Fact[n - k] * Fact[k];
}

mint fact(int n) {return Fact[n];}
mint perm(int n, int k) {
    assert(0 <= n);
    return Fact[n] * iFact[n - k];
}

mint inv2[FACT_SIZE + 1];

} int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  for(int i = 1; i <= FACT_SIZE; i++) {
    mint x = iFact[i] * Fact[i - 1];
    inv2[i] = x * x;
  }
  int n;
  cin >> n;
  vector<mint> q(n);
  rep(i, n) {
    int x;
    cin >> x;
    q[i] = x;
  }
  VVI to(n);
  rep(_, n - 1) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    to[u].emplace_back(v);
    to[v].emplace_back(u);
  }
  VI sz(n), par(n);
  vector<mint> ans = q;
  auto dfs = [&](auto&& self, int u) -> void {
    auto get_total_size = [&](auto&& self, int u) -> int {
      static array_set<200000> visited;
      static VI q;
      visited.clear();
      q.emplace_back(u);
      visited.insert(u);
      while(!q.empty()) {
        int u = q.back(); q.pop_back();
        for(int v: to[u]) if (!visited.contains(v)) {
          q.emplace_back(v);
          visited.insert(v);
        }
      }
      return visited.size;
    };
    int sz_tot = get_total_size(get_total_size, u);
    if (sz_tot == 1) {
      return;
    } else if (sz_tot == 2) {
      int v = to[u][0];
      ans[u] += inv2[2] * q[v];
      ans[v] += inv2[2] * q[u];
      return;
    }
    auto get_centroid = [&](auto&& self, int u, int p=-1) -> int {
      sz[u] = 1;
      par[u] = p;
      int res = -1;
      for(int v: to[u]) if (v != p) {
        int r = self(self, v, u);
        if (r >= 0) res = r;
        sz[u] += sz[v];
      }
      if (res == -1 && sz_tot - sz[u] <= sz_tot / 2) {
        bool ok = true;
        for(int v: to[u]) if (v != p) {
          if (sz[v] > sz_tot / 2) {
            ok = false;
            break;
          }
        }
        if (ok) res = u;
      }
      return res;
    };
    int c = get_centroid(get_centroid, u);
    for(int v : to[c]) if (v == par[c]) {
      sz[v] = sz_tot - sz[c];
    }
    // 1/3, 2/3
    VI e1, e2;
    int x = -1;
    for(int v : to[c]) if (sz[v] >= sz_tot / 3) {
      x = v;
      break;
    }
    if (x != -1) {
      e1.emplace_back(x);
      for(int v : to[c]) if (v != x) {
        e2.emplace_back(v);
      }
    } else {
      int now = 0;
      for(int v: to[c]) {
        if (now + sz[v] <= sz_tot / 2) {
          e1.emplace_back(v);
          now += sz[v];
        } else {
          e2.emplace_back(v);
        }
      }
    }
    to[c].clear();
    static array_unordered_map<int, 200000> mp[2];
    int depth_mx[2] = {};
    auto dfs = [&](int u, auto& mp) {
      mp.clear();
      static VI st;
      mp[u] = 0;
      st.emplace_back(u);
      while(!st.empty()) {
        int u = st.back(); st.pop_back();
        for(int v: to[u]) if (!mp.contains(v)) {
          mp[v] = mp[u] + 1;
          st.emplace_back(v);
        }
      }
    };
    rep(z, 2) {
      swap(to[c], e1);
      dfs(c, mp[z]);
      for(auto [_, d] : mp[z]) chmax(depth_mx[z], d);
      swap(to[c], e1);
      swap(e1, e2);
    }
    int len = depth_mx[0] + depth_mx[1] - 1;
    vector<mint> g(len);
    rep(i, len) g[i] = inv2[i + 3];
    rep(z, 2) {
      vector<mint> f(depth_mx[z]);
      for(auto [u, d] : mp[z]) if (d) {
        f[depth_mx[z] - d] += q[u];
      }
      auto h = convolution(f, g);
      for(auto [u, d] : mp[1 - z]) if (d) {
        ans[u] += h[depth_mx[z] - 1 + (d - 1)];
      }
    }
    rep(_, 2) {
      to[c] = move(e1);
      self(self, c);
      swap(e1, e2);
    }
  };
  dfs(dfs, 0);
  mint k0 = fact(n) * fact(n);
  rep(i, n) cout << (k0 * ans[i]).val() << '\n';
}
yukicoder

結果

ソースコード
