Submission #1045510

#	Time	Username	Problem	Language	Result	Execution time	Memory
1045510		HunterXD	Mergers (JOI19_mergers)	C++17	0 / 100	119 ms	262144 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);
  }
}

// similar to union find path compression
int last_stop = 0;
int path_compression(int u, int mn){
  if(level[u] == mn){
    last_stop = u;
    return parent[u];
  }
  state_group.join(s[u], s[parent[u]]);
  return parent[u] = path_compression(parent[u], mn);
}

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<int> st, st_next;

    for(int u: state[i]) mn = min(mn, level[u]);
    
    // LCA + Path compression
    for(int u: state[i]){
      path_compression(u, mn);
      st.insert(last_stop);
    } 

    while(st.size() > 1){
      for(int u: st){
        path_compression(u, level[u]-1);
        st_next.insert(last_stop);
      }
      st = st_next;
      st_next.clear();
    }
  }

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

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