Submission #1192406

#	Time	Username	Problem	Language	Result	Execution time	Memory
1192406		mannshah1211	Capital City (JOI20_capital_city)	C++20	100 / 100	842 ms	49404 KiB

capital_city

/**
 *   author: tourist
 *   created: 27.04.2025 12:05:50
**/
#include <bits/stdc++.h>

using namespace std;

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

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int n, k;
  cin >> n >> k;
  vector<vector<int>> g(n);
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    --u; --v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  vector<int> c(n);
  vector<vector<int>> at(k);
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    --c[i];
    at[c[i]].push_back(i);
  }
  vector<bool> alive(n, true);
  vector<int> sub(n), freq(k), pv(n), depth(n);
  set<int> used, colors, done;
  vector<int> cur;
  int ans = n + 1, sz = n;
  auto Dfs = [&](auto&& self, int v, int pr) -> void {
    sub[v] = 1;
    freq[c[v]] += 1;
    used.insert(c[v]);
    cur.push_back(v);
    for (int u : g[v]) {
      if (u != pr && alive[u]) {
        self(self, u, v);
        sub[v] += sub[u];
      }
    }
  };
  auto Valid = [&](int col) {
    return (freq[col] == at[col].size());
  };
  auto Find = [&](auto&& self, int v, int pr) -> int {
    for (int u : g[v]) {
      if (u != pr && alive[u] && 2 * sub[u] > sz) {
        return self(self, u, v);
      }
    }
    return v;
  };
  auto Depths = [&](auto&& self, int v, int pr) -> void {
    pv[v] = pr;
    for (int u : g[v]) {
      if (u != pr && alive[u]) {
        depth[u] = depth[v] + 1;
        self(self, u, v);
      }
    }
  };
  auto Add = [&](int col) {
    if (colors.find(col) == colors.end() && done.find(col) == done.end()) {
      colors.insert(col);
    }
  };
  vector<int> p(n), root(n);
  iota(root.begin(), root.end(), 0);
  iota(p.begin(), p.end(), 0);
  int LINES_PRINTED = 0;
  auto Get = [&](auto&& self, int v) -> int {
    if (p[v] == v) {
      return v;
    }
    return p[v] = self(self, p[v]);
  };
  auto Root = [&](int v) {
    return root[Get(Get, v)];
  };
  auto Unite = [&](int u, int v) {
    u = Get(Get, u); v = Get(Get, v);
    if (u == v) {
      return;
    }
    if (rand() & 1) {
      swap(u, v);
    }
    p[v] = u;
    debug(depth[root[v]], depth[root[u]]);
    if (depth[root[v]] < depth[root[u]]) {
      root[u] = root[v];
    }
  };
  auto Reset = [&](int v) {
    p[v] = root[v] = v;
  };
  auto Solve = [&](auto&& self, int v, int pr) -> void {
    Dfs(Dfs, v, -1);
    sz = sub[v];
    int centroid = Find(Find, v, -1);
    depth[centroid] = 0;
    colors.insert(c[centroid]);
    Depths(Depths, centroid, -1);
    bool ok = true;
    while (!colors.empty()) {
      int col = *colors.begin();
      done.insert(col);
      colors.erase(col);
      if (!Valid(col)) {
        ok = false;
        break;
      }
      for (int u : at[col]) {
        while (u != centroid) {
          u = Root(u);
          if (u == centroid) {
            break;
          }
          Add(c[pv[u]]);
          Unite(pv[u], u);
          u = pv[u];
        }
      }
    }
    alive[centroid] = false;
    if (ok) {
      ans = min(ans, static_cast<int>(done.size()));
    }
    for (int col : used) {
      freq[col] = 0;
    }
    for (int v : cur) {
      Reset(v);
    }
    done.clear();
    colors.clear();
    used.clear();
    cur.clear();
    for (int u : g[centroid]) {
      if (alive[u]) {
        self(self, u, centroid);
      }
    }
  };
  Solve(Solve, 0, -1);
  cout << ans - 1 << '\n';
  return 0;
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59