Submission #520873

#	Time	Username	Problem	Language	Result	Execution time	Memory
520873		cig32	Mergers (JOI19_mergers)	C++17	0 / 100	3055 ms	30148 KiB

mergers

#include "bits/stdc++.h"
#define int long long
using namespace std;
const int MAXN = 5e5 + 10;
vector<int> adj[MAXN];
vector<int> adj2[MAXN];
int s[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);
}
int dfs(int node, int prev, int q, int stat) {
  //stat: who the ancestors contain color q
  int cnt = (s[node] == q);
  stat |= (s[node] == q);
  vector<int> v;
  for(int x: adj[node]) {
    if(x != prev) {
      int d = dfs(x, node, q, stat);
      cnt += d;
      if(d > 0 && stat) {
        v.push_back(x);
      }
    }
  }
  for(int x: v) {
    if(set_of(node) != set_of(x)) {
      union_(node, x);
    }
  }
  //ret: how many color q in subtree
  return cnt;
}
int32_t main() {
  int n,k;cin >> n >> k;
  for(int i=1;i<n;i++) {
    int a,b;cin >> a >> b;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  for(int i=1; i<=n; i++) cin >> s[i];
  for(int i=0; i<=n; i++) dsu[i] = i;
  for(int i=1; i<=k; i++) {
    for(int j=1; j<=n; j++) {
      if(s[j] == i) {
        dfs(j, -1, i, 0);
        break;
      }
    }
  }
  map<pair<int, int>, bool> mp;
  for(int i=1; i<=n; i++) {
    for(int x: adj[i]) {
      if(set_of(i) != set_of(x) && mp[{set_of(i), set_of(x)}] == 0) {
        mp[{set_of(i), set_of(x)}] = mp[{set_of(x), set_of(i)}] = 1;
        adj2[set_of(i)].push_back(set_of(x));
        adj2[set_of(x)].push_back(set_of(i));
      }
    }
  }
  int s = 0;
  for(int i=1; i<=n; i++) {
    s += (adj2[i].size() == 1);
  }
  cout << max(0ll, s - 1) << "\n";
}

/*
11 6
1 2
2 3
4 5
5 3
6 7
7 3
8 9
9 3
10 11
11 3
2 2 6 1 1 5 5 4 4 3 3

11 2
1 2
2 3
4 5
5 3
6 7
7 3
8 9
9 3
10 11
11 3
1 1 1 2 1 2 1 2 1 2 1

10 5
1 2
2 3
3 4
4 5
5 6
5 7
7 8
8 9
9 10
1 2 1 2 5 4 3 4 3 4
*/

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