Submission #366398

#	Time	Username	Problem	Language	Result	Execution time	Memory
366398		Kazalika	Mergers (JOI19_mergers)	C++14	100 / 100	1573 ms	226828 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
 
const int N = 5e5 + 5;
 
int n, k, s[N], cnt[N];
 
vector<int> g[N];
 
int ptr[N];
map<int, int> mp[N];
int comp[N], deg[N];
 
int gptr(int v) {
  if (v == ptr[v])
    return v;
  return ptr[v] = gptr(ptr[v]);
}
 
void merge(int a, int b) {
  a = gptr(a), b = gptr(b);
  if (mp[a].size() > mp[b].size())
    swap(a, b);
  ptr[a] = b;
  for (auto &p : mp[a]) {
    mp[b][p.first] += p.second;
    if (mp[b][p.first] == cnt[p.first])
      comp[b]++;
  }
}
 
vector<pair<int, int>> gd;
 
void go(int v, int p = -1) {
  ptr[v] = v;
  mp[ptr[v]][s[v]]++;
  if (cnt[s[v]] == 1)
    comp[ptr[v]]++;
  for (int u : g[v]) {
    if (u == p)
      continue;
    go(u, v);
    if (comp[gptr(u)] == mp[gptr(u)].size()) {
      deg[v]++;
      deg[u]++;
    } else
      gd.emplace_back(v, u);
    merge(v, u);
  }
}
 
int par[N], sm[N], rnk[N];
 
int gpar(int a) {
  if (a == par[a])
    return a;
  return par[a] = gpar(par[a]);
}
void mrg(int a, int b) {
  a = gpar(a), b = gpar(b);
  if (a == b)
    return;
  if (rnk[a] > rnk[b])
    swap(a, b);
  rnk[b]++;
  par[a] = b;
}
 
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cin >> n >> k;
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  for (int i = 0; i < n; ++i) {
    cin >> s[i];
    cnt[s[i]]++;
  }
  go(0);
  iota(par, par + n, 0);
  for (auto &e : gd)
    mrg(e.first, e.second);
  for (int i = 0; i < n; ++i)
    sm[gpar(i)] += deg[i];
  int res = 0;
  for (int i = 0; i < n; ++i)
    if (sm[i] == 1)
      res++;
  cout << (res + 1) / 2;
}

Compilation message (stderr)

mergers.cpp: In function 'void go(int, int)':
mergers.cpp:43:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::map<int, int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   43 |     if (comp[gptr(u)] == mp[gptr(u)].size()) {
      |         ~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~

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