Submission #1257118

#TimeUsernameProblemLanguageResultExecution timeMemory
1257118LucaLucaMCapital City (JOI20_capital_city)C++20
100 / 100
313 ms30164 KiB
#include <iostream>
#include <vector>
#include <algorithm>
#include <cassert>
#include <set>

using ll = long long;
#define debug(x) #x << " = " << x << '\n'

const int INF = 1e9;

int main() {
  std::ios_base::sync_with_stdio(false);
  std::cin.tie(0);

  int n, k;
  std::cin >> n >> k;

  std::vector<std::vector<int>> g(n);

  for (int i = 1; i < n; i++) {
    int u, v;
    std::cin >> u >> v;
    u--, v--;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  std::vector<int> c(n);
  for (int &x : c)  {
    std::cin >> x;
    x--;
  }

  std::vector<bool> blocked(n, false);
  std::vector<int> sz(n);
  std::vector<int> f(k, 0);
  std::vector<bool> vis(n, false);
  std::vector<bool> blockedcolor(k, false);
  std::vector<int> vertices;
  std::vector<int> parent(n);

  std::vector<std::vector<int>> thisColor(k);
  for (int i = 0; i < n; i++) {
    thisColor[c[i]].push_back(i);
  }


  auto calc_sizes = [&](auto &&self, int u, int p) -> void {
    sz[u] = 1;
    for (int v : g[u]) {
      if (v != p && !blocked[v]) {
        self(self, v, u);
        sz[u] += sz[v];
      }
    }
  };

  auto find_centroid = [&](auto &&self, int u, int p, int tsz) -> int {
    for (int v : g[u]) {
      if (v != p && !blocked[v] && sz[v] * 2 > tsz) { 
        return self(self, v, u, tsz);
      } 
    }
    return u;
  };  

  auto calc_dfs = [&](auto &&self, int u) -> void {
    vertices.push_back(u);
    for (int v : g[u]) {
      if (v != parent[u] && !blocked[v]) {
        parent[v] = u;
        self(self, v);
      }
    }
  };

  int answer = k - 1;
  auto build_decomp = [&](auto &&self, int u) -> void {
    calc_sizes(calc_sizes, u, -1);
    int centroid = find_centroid(find_centroid, u, -1, sz[u]);
    parent[centroid] = -1;
    calc_dfs(calc_dfs, centroid);

    if (!blockedcolor[c[centroid]]) {
      std::vector<int> mark;
      std::vector<int> unmark;
      for (int u : vertices) {
        f[c[u]]++;
      }
      int cur = 0;
      if (f[c[centroid]] < (int) thisColor[c[centroid]].size()) {
        cur = 1e9;
      } else {
        for (int u : thisColor[c[centroid]]) {
          mark.push_back(u);
          unmark.push_back(u);
          vis[u] = true;
        }
      }
      while (!mark.empty() && cur < answer) {
        int u = mark.back();
        mark.pop_back();
        if (parent[u] != -1 && !vis[parent[u]]) {
          cur++;
          int newC = c[parent[u]];
          if (f[newC] < (int) thisColor[newC].size()) {
            cur = 1e9;
          } else {
            for (int v : thisColor[newC]) {
              vis[v] = true;
              mark.push_back(v);
              unmark.push_back(v);
            }
          }
        }
      }
      for (int u : unmark) {
        vis[u] = false;
      }
      answer = std::min(answer, cur);
      for (int u : vertices) {
        f[c[u]]--;
      }
    }
    
    blocked[centroid] = true;
    blockedcolor[c[centroid]] = true;
    vertices.clear();
    for (int v : g[centroid]) {
      if (!blocked[v]) {
        self(self, v);
      }
    }
  };

  build_decomp(build_decomp, 0);

  std::cout << answer;

  return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...