Submission #1042974

#	Time	Username	Problem	Language	Result	Execution time	Memory
1042974		juicy	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	238 ms	41652 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;

#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif

const int N = 2e5 + 5;

int n, m;
int c[N], dep[N], res[N], sf[N], mx[N], l[N], lst[N], pos[N];
vector<int> g[N];

void dfs(int u, int p) {
  mx[u] = u;
  for (int v : g[u]) {
    if (v != p) {
      dep[v] = dep[u] + 1;
      dfs(v, u);
      if (dep[mx[v]] > dep[mx[u]]) {
        mx[u] = mx[v];
      }
    }
  }
}

int top;
int st[N];
vector<array<int, 2>> his;

void pop(int k) {
  int l = 1, r = top, p = 0;
  while (l <= r) {
    int md = (l + r) / 2;
    if (dep[st[md]] < k) {
      p = md;
      l = md + 1;
    } else {
      r = md - 1;
    }
  }
  his.push_back({top, st[p]});
  top = p;
}

void push(int k) {
  his.push_back({top, st[top + 1]});
  st[++top] = k;
  pos[k] = top;
}

void roll() {
  assert(his.size());
  auto [a, b] = his.back(); his.pop_back();
  st[top] = b;
  top = a;
}

bool check(int k) {
  return pos[k] <= top && st[pos[k]] == k;
}

void DFS(int u, int p) {
  int prv = -1;
  pop(dep[u] - l[u]);
  if (u ^ p) {
    if (!lst[c[p]] || !check(lst[c[p]])) {
      prv = lst[c[p]];
      lst[c[p]] = p;
      push(p);
    }
  }
  int h = dep[mx[u]] - dep[u];
  pop(dep[u] - h);
  res[u] = max(res[u], top);
  roll();
  {
    int k = g[u].size();
    sf[k] = 0;
    for (int i = k - 1; i >= 0; --i) {
      sf[i] = sf[i + 1];
      int v = g[u][i];
      if (v != p) {
        sf[i] = max(sf[i], dep[mx[v]] - dep[u] + 1);
      }
    }
    int pf = 0;
    for (int i = 0; i < k; ++i) {
      int v = g[u][i];
      if (v != p) {
        l[v] = max(pf, sf[i + 1]);
        pf = max(pf, dep[mx[v]] - dep[u] + 1);
      }
    }
    for (int v : g[u]) {
      if (v != p) {
        DFS(v, u);
      }
    }
  }
  if (~prv) {
    lst[c[p]] = prv;
    roll();
  }
  roll();
  l[u] = 0;
}

void solve(int u) {
  dep[u] = 1;
  dfs(u, u);
  DFS(u, u);
}

int findbest(int r) {
  dep[r] = 0;
  dfs(r, r);
  return mx[r];
}

array<int, 2> diam() {
  int u = findbest(1);
  return {u, findbest(u)};
}

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

  cin >> n >> m;
  for (int i = 1; i < n; ++i) {
    int u, v; cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }  
  for (int i = 1; i <= n; ++i) {
    cin >> c[i];
  }
  auto [a, b] = diam();
  solve(a);
  solve(b);
  for (int i = 1; i <= n; ++i) {
    cout << res[i] << "\n";
  }
  return 0;
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32