Submission #711483

#	Time	Username	Problem	Language	Result	Execution time	Memory
711483		Foxyy	Mergers (JOI19_mergers)	C++17	100 / 100	959 ms	196524 KiB

mergers

#include <bits/stdc++.h>
using namespace std;

#define Foxyy cin.tie(0); cout.sync_with_stdio(0);
#define ll long long

namespace HLD {
	int n;
	vector<vector<int>>& adj = *(new vector<vector<int>>());
	vector<int> chainRootOfNode{};
	vector<int> heavyChild{};
	vector<int> parent{};
	vector<int> height{};
	vector<int> depth{};
	
	void scanHeavy(int u, int p) {
		if (p != -1) {
			depth[u] = depth[p] + 1;
		}
		parent[u] = p;
		int currentHeavyChild = -1;
		for (int v : adj[u]) if (v != p) {
			scanHeavy(v, u);
			if (currentHeavyChild == -1 or height[v] + 1 > height[currentHeavyChild]) {
				currentHeavyChild = v;
			}
		}
		heavyChild[u] = currentHeavyChild;
		height[u] = height[heavyChild[u]] + 1;
	}
	
	void buildChains(int u, int root) {
		chainRootOfNode[u] = root;
		for (int v : adj[u]) if (v != parent[u]) {
			if (v == heavyChild[u]) {
				buildChains(v, root);
			} else {
				buildChains(v, v);
			}
		}
	}

	void initialize(vector<vector<int>>& _adj) {
		n = (int)_adj.size();
		adj = _adj;
		chainRootOfNode.resize(n, -1);
		parent.resize(n);
		heavyChild.resize(n, -1);
		height.resize(n);
		depth.resize(n);
		
		chainRootOfNode[0] = parent[0] = 0;
		depth[0] = 0;
		
		scanHeavy(0, 0);
		buildChains(0, 0);
	}
	
	int getLCAOf(int a, int b) {
		while (chainRootOfNode[a] != chainRootOfNode[b]) {
			if (depth[chainRootOfNode[a]] < depth[chainRootOfNode[b]]) {
				swap(a, b);
			}
			a = parent[chainRootOfNode[a]];
		}
		if (depth[a] > depth[b]) {
			return b;
		} else {
			return a;
		}
	}
	
	int getChainRootOfNode(int u) {
		return chainRootOfNode[u];
	}
}; // namespace HLD

struct Solver {
	int n;
	int k;
	vector<vector<int>>& adj;
	vector<int>& s;
	
	vector<int> f;
	
	int get(int x) {
		return x == f[x] ? x : f[x] = get(f[x]);
	}
	
	void solve() {
		HLD::initialize(adj);
//		cerr << "init\n";
		vector<vector<int>> stateCities(k);
		for (int i = 0; i < n; i++) {
			stateCities[s[i]].push_back(i);
		}
		
		f.resize(n);
		iota(f.begin(), f.end(), 0);
		
		for (int i = 0; i < k; i++) if (not stateCities[i].empty()) {
			int u = stateCities[i][0];
			for (int v : stateCities[i]) {
//				cerr << u << ' ' << v << '\n';
				u = get(u), v = get(v);
				int lca = get(HLD::getLCAOf(u, v));
//				cerr << i << ' ' << u << ' ' << v << ' ' << lca << '\n';
				while (u != lca) {
//					cerr << u << '\n';
					f[u] = lca;
					u = get(HLD::parent[u]);
				}
				while (v != lca) {
					f[v] = lca;
//					cerr << v << '\n';
					v = get(HLD::parent[v]);
				}
			}
		}
		
		vector<set<int>> s(n);
		for (int i = 0; i < n; i++) {
			for (int j : adj[i]) if (get(i) != get(j)) {
				s[get(i)].insert(get(j));
			}
		}
		
		int cnt = 0;
		for (int i = 0; i < n; i++) {
			if (s[i].size() == 1u) {
				cnt++;
			}
		} 	
		
//		cerr << cnt << '\n';
		
		cout << (cnt+1)/2 << '\n';
	}
};

signed main() {
	Foxyy
	
	int T = 1;
//	cin >> T;
	while (T--) {
		int n, k;
		cin >> n >> k;
		vector<vector<int>> adj(n);
		vector<int> s(n);
		for (int i = 0; i < n-1; i++) {
			int a, b;
			cin >> a >> b;
			a--, b--;
			adj[a].push_back(b);
			adj[b].push_back(a);
		}
		for (int i = 0; i < n; i++) {
			cin >> s[i];
			s[i]--;
		}
		
		Solver solver{n, k, adj, s};
		solver.solve();
	}
}

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