Submission #925716

#	Time	Username	Problem	Language	Result	Execution time	Memory
925716		cig32	Capital City (JOI20_capital_city)	C++17	100 / 100	381 ms	134668 KiB

capital_city

#include "bits/stdc++.h"
using namespace std;
#define int long long
#define double long double
const int MAXN = 5e5 + 10;
const int MOD = 1e9 + 7;
mt19937_64 rng((int)std::chrono::steady_clock::now().time_since_epoch().count());
int rnd(int x, int y) {
  int u = uniform_int_distribution<int>(x, y)(rng); return u;
}
int bm(int b, int p) {
  if(p==0) return 1 % MOD;
  int r = bm(b, p >> 1);
  if(p&1) return (((r*r) % MOD) * b) % MOD;
  return (r*r) % MOD;
}
int inv(int b) { 
  return bm(b, MOD-2);
}
int fastlog(int x) {
  return (x == 0 ? -1 : 64 - __builtin_clzll(x) - 1);
}
void printcase(int i) { cout << "Case #" << i << ": "; }
static void run_with_stack_size(void (*func)(void), size_t stsize) {
  char *stack, *send;
  stack = (char *)malloc(stsize);
  send = stack + stsize - 16;
  send = (char *)((uintptr_t)send / 16 * 16);
  asm volatile(
    "mov %%rsp, (%0)\n"
    "mov %0, %%rsp\n"
    :
    : "r"(send));
  func();
  asm volatile("mov (%0), %%rsp\n" : : "r"(send));
  free(stack);
}
int n, k;
int c[MAXN], sz[MAXN];
vector<int> tr[MAXN];
vector<int> adj[MAXN], rev[MAXN];
stack<int> ts;
bool vis[MAXN];
void dfs1(int node) {
  vis[node] = 1;
  for(int x: adj[node]) {
    if(!vis[x]) dfs1(x);
  }
  ts.push(node);
}
vector<int> comp;
vector<int> cond[MAXN];
int dsu[MAXN];
int set_of(int u) {
  if(dsu[u] == u) return u;
  return dsu[u] = set_of(dsu[u]);
}
void union_(int u, int v) {
  dsu[set_of(u)] = set_of(v);
}
void dfs2(int node) {
  vis[node] = 1;
  for(int x: rev[node]) {
    if(!vis[x]) {
      dfs2(x);
    }
  }
  comp.push_back(node);
}
int tin[MAXN], cur;
void dfs0(int node, int prv) {
  tin[node] = ++cur;
  sz[node] = 1;
  for(int x: tr[node]) {
    if(x != prv) {
      dfs0(x, node);
      sz[node] += sz[x];
    }
  }
}
set<int> city[MAXN];
bool have(int node, int prv, int tar) {
  if(tin[node] < tin[prv]) {
    if(*city[tar].rbegin() >= tin[prv] + sz[prv]) return 1;
    if(*city[tar].begin() < tin[prv]) return 1;
    return 0;
  }
  else {
    // [tin[node], tin[node] + sz[node])
    auto it = city[tar].lower_bound(tin[node]);
    if(it != city[tar].end() && (*it) < tin[node] + sz[node]) return 1;
    return 0;
  }
  return 0;
}
void solve(int tc) { 
  cin >> n >> k;
  vector<pair<int,int> > edges;
  for(int i=1; i<n; i++) {
    int u, v;
    cin >> u >> v;
    edges.push_back({u, v});
    tr[u].push_back(v);
    tr[v].push_back(u);
  }
  for(int i=1; i<=n; i++) {
    cin >> c[i];
  }
  dfs0(1, -1);
  for(int i=1; i<=n; i++) {
    city[c[i]].insert(tin[i]);
  }
  for(pair<int, int> k : edges) {
    if(have(k.first, k.second, c[k.second])) {
      adj[c[k.second]].push_back(c[k.first]);
    }
    if(have(k.second, k.first, c[k.first])) {
      adj[c[k.first]].push_back(c[k.second]);
    }
  }
  for(int i=1; i<=k; i++) {
    for(int x: adj[i]) {
      rev[x].push_back(i);
    }
  }
  for(int i=1; i<=k; i++) {
    if(!vis[i]) {
      dfs1(i);
    }
  }
  for(int i=1; i<=k; i++) vis[i] = 0, dsu[i] = i;
  while(ts.size()) {
    int tp = ts.top();
    ts.pop();
    if(!vis[tp]) {
      comp.clear();
      dfs2(tp);
      for(int i=1; i<comp.size(); i++) {
        union_(comp[i-1], comp[i]);
      }
    }
  }
  set<pair<int,int>> newedges;
  for(int i=1; i<=k; i++) {
    for(int x: adj[i]) {
      if(set_of(i) != set_of(x) && !newedges.count({set_of(i), set_of(x)})) {
        cond[set_of(i)].push_back(set_of(x));
        newedges.insert({set_of(i), set_of(x)});
      }
    }
  }
  int ans = 1e9;
  int cnt[k+1];
  for(int i=1; i<=k; i++) {
    cnt[i] = 0;
  }
  for(int i=1; i<=k; i++) {
    cnt[set_of(i)]++;
  }

  for(int i=1; i<=k; i++) {
    if(set_of(i) == i && cond[i].empty()) {
      ans = min(ans, cnt[i]);
    }
  }
  cout << ans - 1 << "\n";
}

void uwu() {
  ios::sync_with_stdio(0); cin.tie(0);
  int t = 1; //cin >> t;
  for(int i=1; i<=t; i++) solve(i);
}
int32_t main() {
  #ifdef ONLINE_JUDGE
  uwu();
  #endif
  #ifndef ONLINE_JUDGE
  run_with_stack_size(uwu, 1024 * 1024 * 1024); // run with a 1 GiB stack
  #endif
}
/*
g++ A.cpp -std=c++17 -O2 -o A
./A < input.txt
*/

Compilation message (stderr)

capital_city.cpp: In function 'void solve(long long int)':
capital_city.cpp:138:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  138 |       for(int i=1; i<comp.size(); i++) {
      |                    ~^~~~~~~~~~~~

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