제출 #137971

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
137971		Bohoty	Mergers (JOI19_mergers)	C++14	0 / 100	205 ms	20948 KiB

mergers

#pragma GCC optimize ("O3")
#include<bits/stdc++.h>
using namespace std;
#define pb push_back
#define F first
#define S second
using ll = long long;
using ii = pair<int, int>;
const int N = 5e5 + 5, mod = 1e9 + 7;
int n, k;
std::vector<int> adj[N];
int in[N], out[N], ord[N], lvl[N], sz[N], timer;
int cnt[N];
int cur[N];
int col[N];
int freq[N];
void dfs0(int u){
  sz[u] = 1;
  in[u] = timer;
  for(auto v : adj[u]){
    lvl[v] = lvl[u] + 1;
    adj[v].erase(find(adj[v].begin(), adj[v].end(), u));
    dfs0(v);
    sz[u] += sz[v];
  }
  ord[timer] = u;
  out[u] = timer++;
}
int dfs1(int u, bool keep) {
  int mx = -1, bigChild = -1;
  for(auto v : adj[u])if(sz[v] > mx){
    mx = sz[v], bigChild = v;
  }
  if(bigChild == -1){
    cur[u] = k - 1;
    cnt[col[u]]++;
    cur[u] += (cnt[col[u]] == freq[col[u]]);
    int ret = cur[u] == k;
    // cout << freq[col[u]] << ' ' << cnt[col[u]] << ' ' << cur[u] << ' ' << ret  << '\n';
    if(!keep){
      cur[u] -= (cnt[col[u]] == freq[col[u]]);
      cnt[col[u]]--;
      cur[u] += (cnt[col[u]] == 0);
    }
    // cout << u << ' ' << ret << '\n';
    return ret;
  }
  int ret = 0;
  for(auto v : adj[u])if(v != bigChild){
    ret += dfs1(v, 0);
  }
  // cout << u << ' ' << ret << '\n';
  ret += dfs1(bigChild, 1);
  cur[u] = cur[bigChild];
  for(auto v : adj[u])
  if(v != bigChild){
    for(int x = in[v]; x <= out[v]; x++){
      cur[u] -= (cnt[col[ord[x]]] == 0);
      cnt[col[ord[x]]]++;
      cur[u] += (cnt[col[ord[x]]] == freq[col[ord[x]]]);
    }
  }
  cur[u] -= (cnt[col[u]] == 0);
  cnt[col[u]]++;
  cur[u] += (cnt[col[u]] == freq[col[u]]);
  if(!ret && u != 1){
    ret = cur[u] == k;
  }
  if(!keep){
    for(int x = in[u]; x <= out[u]; x++){
      cur[u] -= (cnt[col[ord[x]]] == freq[col[ord[x]]]);
      cnt[col[ord[x]]]--;
      cur[u] += (cnt[col[ord[x]]] == 0);
    }
  }
  return ret;
}
int main(){

  cin >> n >> k;

  for(int i = 0, u, v; i < n - 1; i++){
    cin >> u >> v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  for(int i = 1; i <= n; i++){
    cin >> col[i];
    freq[col[i]]++;
  }
  dfs0(1);


  cout << (dfs1(1,0) + 1) / 2 << '\n';

  return 0;
}

• Subtask 1: 0 / 10
• Subtask 2: 0 / 24
• Subtask 3: 0 / 14
• Subtask 4: 0 / 22
• Subtask 5: 0 / 30