結果

問題 No.2634 Tree Distance 3
コンテスト
ユーザー emthrm
提出日時 2024-02-16 23:47:31
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
TLE  
実行時間 -
コード長 6,149 bytes
コンパイル時間 4,019 ms
コンパイル使用メモリ 290,196 KB
実行使用メモリ 58,472 KB
最終ジャッジ日時 2024-09-28 22:52:56
合計ジャッジ時間 14,901 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample -- * 2
other TLE * 1 -- * 68
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;
#define FOR(i,m,n) for(int i=(m);i<(n);++i)
#define REP(i,n) FOR(i,0,n)
using ll = long long;
constexpr int INF = 0x3f3f3f3f;
constexpr long long LINF = 0x3f3f3f3f3f3f3f3fLL;
constexpr double EPS = 1e-8;
constexpr int MOD = 998244353;
// constexpr int MOD = 1000000007;
constexpr int DY4[]{1, 0, -1, 0}, DX4[]{0, -1, 0, 1};
constexpr int DY8[]{1, 1, 0, -1, -1, -1, 0, 1};
constexpr int DX8[]{0, -1, -1, -1, 0, 1, 1, 1};
template <typename T, typename U>
inline bool chmax(T& a, U b) { return a < b ? (a = b, true) : false; }
template <typename T, typename U>
inline bool chmin(T& a, U b) { return a > b ? (a = b, true) : false; }
struct IOSetup {
  IOSetup() {
    std::cin.tie(nullptr);
    std::ios_base::sync_with_stdio(false);
    std::cout << fixed << setprecision(20);
  }
} iosetup;

struct LowestCommonAncestorByDoubling {
  std::vector<int> depth;

  explicit LowestCommonAncestorByDoubling(
      const std::vector<std::vector<int>>& graph)
      : is_built(false), n(graph.size()),
        table_h(std::countr_zero(std::bit_floor(graph.size())) + 1),
        graph(graph) {
    depth.resize(n);
    parent.resize(table_h, std::vector<int>(n));
  }

  void build(int root = 0) {
    is_built = true;
    dfs(-1, root, 0);
    for (int i = 0; i + 1 < table_h; ++i) {
      for (int ver = 0; ver < n; ++ver) {
        parent[i + 1][ver] =
            (parent[i][ver] == -1 ? -1 : parent[i][parent[i][ver]]);
      }
    }
  }

  int query(int u, int v) const {
    assert(is_built);
    if (depth[u] > depth[v]) std::swap(u, v);
    for (int i = 0; i < table_h; ++i) {
      if ((depth[v] - depth[u]) >> i & 1) v = parent[i][v];
    }
    if (u == v) return u;
    for (int i = table_h - 1; i >= 0; --i) {
      if (parent[i][u] != parent[i][v]) {
        u = parent[i][u];
        v = parent[i][v];
      }
    }
    return parent.front()[u];
  }

  int distance(const int u, const int v) const {
    assert(is_built);
    return depth[u] + depth[v] - depth[query(u, v)] * 2;
  }

  int level_ancestor(int v, const int d) const {
    assert(is_built);
    if (depth[v] < d) return -1;
    for (int i = depth[v] - d, bit = 0; i > 0; i >>= 1, ++bit) {
      if (i & 1) v = parent[bit][v];
    }
    return v;
  }

  int jump(const int u, const int v, const int d) const {
    assert(is_built);
    if (d == 0) [[unlikely]] return u;
    const int l = query(u, v), d_lu = depth[u] - depth[l];
    if (d_lu == d) return l;
    if (d_lu > d) return level_ancestor(u, depth[u] - d);
    return level_ancestor(v, depth[l] + (d - d_lu));
  }

 private:
  bool is_built;
  const int n, table_h;
  const std::vector<std::vector<int>> graph;
  std::vector<std::vector<int>> parent;

  void dfs(const int par, const int ver, const int cur_depth) {
    depth[ver] = cur_depth;
    parent.front()[ver] = par;
    for (const int e : graph[ver]) {
      if (e != par) dfs(ver, e, cur_depth + 1);
    }
  }
};

int main() {
  constexpr int S = 448;
  int n; cin >> n;
  vector<int> a(n);
  for (int& a_i : a) cin >> a_i;
  vector<vector<int>> graph(n);
  REP(_, n - 1) {
    int u, v; cin >> u >> v; --u; --v;
    graph[u].emplace_back(v);
    graph[v].emplace_back(u);
  }
  vector<int> order(n);
  iota(order.begin(), order.end(), 0);
  ranges::sort(order, greater<int>(), [&](const int i) -> int { return a[i]; });
  LowestCommonAncestorByDoubling lca(graph);
  lca.build();
  map<int, vector<int>> bad;
  vector<int> dp(n, 0), ans(n, 0), is_on(n, false);
  for (int i = 0; i < n;) {
    const int j = min(i + S, n);
    if (a[order[j - 1]] == a[order[j]]) bad[a[order[j]]].emplace_back();
    FOR(r, i, j) ans[order[r]] = dp[order[r]];
    FOR(r, i, j) FOR(l, i, r) chmax(ans[order[r]], lca.distance(order[l], order[r]));
    FOR(r, i, j) is_on[order[r]] = true;
    vector<int> dp2(n, -INF);
    const auto dfs1 = [&](auto dfs1, const int par, const int ver) -> void {
      if (is_on[ver]) dp2[ver] = 0;
      for (const int e : graph[ver]) {
        if (e != par) {
          dfs1(dfs1, ver, e);
          chmax(dp2[ver], dp2[e] + 1);
        }
      }
    };
    dfs1(dfs1, -1, 0);
    const auto dfs2 = [&](auto dfs2, const int par, const int ver, const int pdp) -> void {
      dp[ver] = max(dp2[ver], pdp);
      const int c = graph[ver].size();
      vector<int> left(c, -INF), right(c, -INF);
      FOR(i, 1, c) left[i] = max(left[i - 1], graph[ver][i - 1] == par ? pdp : dp2[graph[ver][i - 1]]);
      for (int i = c - 2; i >= 0; --i) {
        right[i] = max(right[i + 1], graph[ver][i + 1] == par ? pdp : dp2[graph[ver][i + 1]]);
      }
      REP(i, c) {
        if (graph[ver][i] != par) dfs2(dfs2, ver, graph[ver][i], max(left[i], right[i]) + 1);
      }
    };
    dfs2(dfs2, -1, 0, -INF);
    i = j;
  }
  REP(i, n) {
    if (const auto it = bad.find(a[i]); it != bad.end()) it->second.emplace_back(i);
  }
  fill(is_on.begin(), is_on.end(), false);
  for (const vector<int>& vs : bad | views::values) {
    fill(dp.begin(), dp.end(), -INF);
    for (const int v : vs) is_on[v] = true;
    vector<int> dp2(n, -INF);
    const auto dfs1 = [&](auto dfs1, const int par, const int ver) -> void {
      if (is_on[ver]) dp2[ver] = 0;
      for (const int e : graph[ver]) {
        if (e != par) {
          dfs1(dfs1, ver, e);
          chmax(dp2[ver], dp2[e] + 1);
        }
      }
    };
    dfs1(dfs1, -1, 0);
    const auto dfs2 = [&](auto dfs2, const int par, const int ver, const int pdp) -> void {
      dp[ver] = max(dp2[ver], pdp);
      const int c = graph[ver].size();
      vector<int> left(c, -INF), right(c, -INF);
      FOR(i, 1, c) left[i] = max(left[i - 1], graph[ver][i - 1] == par ? pdp : dp2[graph[ver][i - 1]]);
      for (int i = c - 2; i >= 0; --i) {
        right[i] = max(right[i + 1], graph[ver][i + 1] == par ? pdp : dp2[graph[ver][i + 1]]);
      }
      REP(i, c) {
        if (graph[ver][i] != par) dfs2(dfs2, ver, graph[ver][i], max(left[i], right[i]) + 1);
      }
    };
    dfs2(dfs2, -1, 0, -INF);
    for (const int v : vs) chmax(ans[v], dp[v]);
    for (const int v : vs) is_on[v] = false;
  }
  REP(i, n) cout << ans[i] << " \n"[i + 1 == n];
  return 0;
}
0