Submission #1045540

#	Time	Username	Problem	Language	Result	Execution time	Memory
1045540		HunterXD	Mergers (JOI19_mergers)	C++17	0 / 100	71 ms	38088 KiB

mergers

#include <bits/stdc++.h>

using namespace std;

#define pb push_back

typedef pair<int, int> pi;

const char nd = '\n';

void solve();

signed main(){
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int t = 1;
  // cin >> t;
  while(t-->0) solve();
}

// ------------------- Code Starts here -------------------

// Union Find
struct UF{
  vector<int> p, s;
  UF(int n): p(n), s(n, 1){
    iota(p.begin(), p.end(), 0);
  }

  int find(int x){
    return p[x] == x ? x : p[x] = find(p[x]);
  }

  void join(int x, int y){
    x = find(x);
    y = find(y);
    if(x == y) return;
    if(s[x] < s[y]) swap(x, y);
    p[y] = x;
    s[x] += s[y];
  }
};


const int N = 5e5+7;

vector<int> g[N], state[N];
int s[N], parent[N], level[N];
UF state_group(N);

// dfs to get parent and level of each node
void dfs(int u, int p){
  parent[u] = p;
  level[u] = level[p]+1;
  for (int v : g[u]){
    if (v == p) continue;
    dfs(v, u);
  }
}

void solve(){
  int n, k;
  cin >> n >> k;

  vector<pi> edges;

  for (int i = 1; i < n; i++){
    int u, v;
    cin >> u >> v;
    edges.pb({u, v});
    g[u].pb(v);
    g[v].pb(u);    
  }

  for (int i = 1; i <= n; i++){
    cin >> s[i];
    state[s[i]].pb(i);
  }

  dfs(1, 0);


  for (int i = 1; i <= n; i++){
    int mn = N;
    set<pi> st;

    for(int u: state[i]) st.insert({-level[u], u});
    
    while(st.size() > 1){
      int v;
      auto [l, u] = *st.begin();
      st.erase(st.begin());

      v = parent[u];
      state_group.join(s[u], s[v]);
      parent[u] = parent[v]; // path compression
      st.insert({-level[v], v});
    }
  }

  // degree of each state group
  vector<int> deg(n+1, 0);

  for(auto [u, v]: edges){
    u = state_group.find(s[u]);
    v = state_group.find(s[v]);
    if (u == v) continue;
    deg[u]++;
    deg[v]++;
  }

  int leaf = 0;
  for(int i = 1; i <= n; i++){
    if(deg[i] == 1) leaf++;
  }

  cout << (leaf + 1) / 2 << nd;
}

Compilation message (stderr)

mergers.cpp: In function 'void solve()':
mergers.cpp:85:9: warning: unused variable 'mn' [-Wunused-variable]
   85 |     int mn = N;
      |         ^~

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