Submission #1042963

#	Time	Username	Problem	Language	Result	Execution time	Memory
1042963		juicy	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	526 ms	41524 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], s[4 * N], lz[4 * N], cnt[4 * N];
vector<int> g[N];

void pul(int id) {
  s[id] = min(s[id * 2], s[id * 2 + 1]);
  cnt[id] = 0;
  for (auto i : {id * 2, id * 2 + 1}) {
    if (s[id] == s[i]) {
      cnt[id] += cnt[i];
    }
  }
}

void bld(int id = 1, int l = 1, int r = n) {
  lz[id] = 0;
  if (l == r) {
    s[id] = 1;
    cnt[id] = 1;
    return;
  }
  int md = (l + r) / 2;
  bld(id * 2, l, md);
  bld(id * 2 + 1, md + 1, r);
  pul(id); 
}

void app(int id, int x) {
  s[id] += x;
  lz[id] += x;
}

void psh(int id) {
  if (lz[id]) {
    app(id * 2, lz[id]);
    app(id * 2 + 1, lz[id]);
    lz[id] = 0;
  }
}

void upd(int u, int v, int x, int id = 1, int l = 1, int r = n) {
  if (u <= l && r <= v) {
    app(id, x);
    return;
  }
  psh(id);
  int md = (l + r) / 2;
  if (u <= md) {
    upd(u, v, x, id * 2, l, md);
  }
  if (md < v) {
    upd(u, v, x, id * 2 + 1, md + 1, r);
  }
  pul(id);
}

int qry(int i, int id = 1, int l = 1, int r = n) {
  if (l == r) {
    return s[id];
  }
  psh(id);
  int md = (l + r) / 2;
  return i <= md ? qry(i, id * 2, l, md) : qry(i, id * 2 + 1, md + 1, r);
}

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];
      }
    }
  }
}

void DFS(int u, int p) {
  int prv = -1;
  if (l[u] && dep[u] > 2) {
    upd(max(1, dep[u] - l[u]), dep[u] - 2, 1);
  }
  if (u ^ p) {
    if (!lst[c[p]] || qry(dep[lst[c[p]]])) {
      prv = lst[c[p]];
      lst[c[p]] = p;
      upd(dep[p], dep[p], -1);
    }
  }
  int h = dep[mx[u]] - dep[u];
  upd(max(1, dep[u] - h), dep[u], 1); 
  if (s[1] == 0) {
    res[u] = max(res[u], cnt[1]);
  }
  upd(max(1, dep[u] - h), dep[u], -1);
  {
    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;
    upd(dep[p], dep[p], 1);
  }
  if (l[u] && dep[u] > 2) {
    upd(max(1, dep[u] - l[u]), dep[u] - 2, -1);
  }
  l[u] = 0;
}

void solve(int u) {
  dep[u] = 1;
  dfs(u, u);
  bld();
  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