Submission #1257091

#	Time	Username	Problem	Language	Result	Execution time	Memory
1257091		LucaLucaM	Capital City (JOI20_capital_city)	C++20	0 / 100	3092 ms	30868 KiB

capital_city

#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);

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

  std::vector<bool> blockedcolor(k, false);
  std::vector<int> parent(n);

  std::vector<int> vertices;

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

  std::vector<bool> vis(n, false);
  std::vector<bool> inSubtree(n, false);

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

  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> st;
      std::vector<int> st2;
      for (int u : vertices) {
        inSubtree[u] = true;
      }
      int need = (int) thisColor[c[centroid]].size();
      int cur = 0;
      for (int u : thisColor[c[centroid]]) {
        if (!inSubtree[u]) {
          cur = 1e9;
          break;
        }
        st.push_back(u);
        vis[u] = true;
      }
      while (!st.empty()) {
        int u = st.back();
        st2.push_back(u);
        st.pop_back();
        if (parent[u] != -1 && !vis[parent[u]]) {
          if (blockedcolor[c[parent[u]]]) {
            cur = 1e9;
            break;
          }
          cur++;
          need += (int) thisColor[c[parent[u]]].size();
          for (int v : thisColor[c[parent[u]]]) {
            if (!inSubtree[v]) {
              cur = 1e9;
              break;
            }
            vis[v] = true;
            st.push_back(v);
          }
        }
      }
      for (int x : st2) {
        vis[x] = false;
      }
      // if (cur == 0 && centroid == 0) {
        // std::cout << "yesSir\n";
        // for (int x : st2) {
        //   std::cout << debug(x);
        // }
      // }
      if (need == (int) st2.size()) {
        // if (cur == 0) {
          // std::cout << " ???\n";
          // for (int x : st2) {
          //   std::cout << debug(x);
          // }
        // }
        answer = std::min(answer, cur);
      }
      for (int u : vertices) {
        inSubtree[u] = false;
      }
    }

    // std::cout << debug(centroid);
    
    blocked[centroid] = true;
    blockedcolor[c[centroid]] = true;
    for (int v : g[centroid]) {
      if (!blocked[v]) {
        self(self, v);
      }
    }
  };

  build_decomp(build_decomp, 0);

  std::cout << answer;

  return 0;
}

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