Submission #1312268

#TimeUsernameProblemLanguageResultExecution timeMemory
1312268vlomaczkCapital City (JOI20_capital_city)C++20
Compilation error
0 ms0 KiB
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
typedef long long ll;
using namespace __gnu_pbds;
using namespace std;

template <typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

template <typename TY>
struct MultiOrderedSet {
	ordered_set<pair<TY, int>> os;
	int cnt = 0;
	void insert(TY x) {
		os.insert({x, cnt++});
	}
	void erase(TY x) {
		int idx = os.order_of_key({x,-1});
		assert(idx < os.size());
		os.erase(*os.find_by_order(idx));
	}
	TY first() { return os.begin()->first; }
	TY last() { return os.rbegin()->first; }
	void clear() {
		while(os.size()) os.erase(*os.begin());
	}
	int size() { return os.size(); }
	bool empty() { return os.empty(); }
	int order_of_key(TY x) {
		return os.order_of_key({x, -1});
	}
	TY find_by_order(ll x) {
		return os.find_by_order(x)->first;
	}
};

struct ds {
	set<int> s;
	map<int, int> mapa;
	MultiOrderedSet<int> mos;
	vector<int> values, S;
	ll dodane = 0;

	void insert(int v) {
		if(s.find(v)==s.end()) {
			s.insert(v);
			values.push_back(-dodane);
			mapa[v] = values.size() - 1;
			S.push_back(v);
		}
	}

	void dodaj1() { dodane++; }
	void odejmij(int x) { mos.insert(x); }
	
	int get_val(ll v) {
		int idx = mapa[v];
		return values[idx] + dodane - (mos.size() - mos.order_of_key(idx));
	}
};

int n,m;
int M = 200'001;
int maxlog = 17;
vector<vector<int>> g(M), j(M, vector<int>(maxlog+1));
vector<int> C(M), sajz(M), path_end(M), d(M), h(M, -1), ile_mialem(M), ans(M), pre(M);
vector<vector<int>> virtuals(M), grab_ans(M), kto_ma(M), gang(M);
vector<vector<vector<int>>> v_childs(M);
vector<ds> struktura(M);
int cnt = 0;

void jump_dfs(int v, int p) {
	pre[v] = ++cnt;
	d[v] = d[p] + 1;
	sajz[v] = 1;
	j[v][0] = p;
	for(int i=1; i<=maxlog; ++i) j[v][i] = j[j[v][i-1]][i-1];
	for(auto u : g[v]) {
		if(u!=p) {
			jump_dfs(u,v);
			sajz[v] += sajz[u];
		}
	}
}

int lca(int a, int b) {
	if(d[b] > d[a]) swap(a,b);
	for(int i=maxlog; i>=0; --i) {
		if(d[j[a][i]] >= d[b]) a = j[a][i];
	}
	if(a==b) return a;
	for(int i=maxlog; i>=0; --i) {
		if(j[a][i]!=j[b][i]) {
			a = j[a][i];
			b = j[b][i];
		}
	}
	return j[a][0];
}

void make_hld(int v, int p) {
	for(auto u : g[v]) if(u!=p && (h[v]==-1 || sajz[u] > sajz[h[v]])) h[v] = u;
	if(h[v]!=-1) {
		path_end[h[v]] = path_end[v];
		make_hld(h[v], v);
	}
	for(auto u : g[v]) {
		if(u==p || u==h[v]) continue;
		path_end[u] = u;
		make_hld(u,v);
	}
}

ds mw_dfs(int v, int p) {
	ds myds;
	if(h[v]!=-1) myds = mw_dfs(h[v], v);
	for(auto u : g[v]) {
		if(u==p || u==h[v]) continue;
		ds cds = mw_dfs(u,v);
		for(auto x : cds.s) myds.insert(x);
	}
	myds.insert(C[v]);
	myds.dodaj1();
	cout << v << ": "; for(auto x : myds.S) cout << "(" << x << " " << myds.get_val(x)<< ") "; cout << "\n";
	int idx = 0;
	for(auto x : virtuals[v]) {
		if(v==2) cout << x << ": ";
		for(int i=0; i<v_childs[v][idx].size()-1; ++i) {
			
			if(d[path_end[u]] <= d[v]) {
				myds.odejmij(ile_mialem[u]);
				if(v==2) cout << "[" << v << " : " << ile_mialem[u] << "] - ";
			} else {
				struktura[path_end[u]].odejmij(ile_mialem[u]);
				if(v==2) cout << "[" << path_end[u] << " : " << ile_mialem[u] << "] - ";
			}
			if(v==2) cout << u << ", ";
		}
		cout << "\n";
		++idx;
	}
	for(auto c : grab_ans[v]) {
		ans[c] = myds.get_val(c);
		for(auto x : kto_ma[c]) {
			if(c==2) cout << x << " - " << struktura[x].get_val(c) << "\n";
			ans[c] += struktura[x].get_val(c);
		}
	}
	for(auto x : myds.S) cout << "(" << x << " " << myds.get_val(x)<< ") "; cout << "\n"; cout << "\n";
	ile_mialem[v] = myds.s.size();
	if(path_end[v] == v) {
		struktura[v] = myds;
		for(auto x : myds.s) kto_ma[x].push_back(v);
	}
	return myds;
}

void make_vt(int c) {
	sort(gang[c].begin(), gang[c].end(), [&](int a, int b){
		return pre[a] < pre[b];
	});
	set<int> verts;
	for(int i=0; i<gang[c].size()-1; ++i) {
		verts.insert(gang[c][i]);
		verts.insert(gang[c][i+1]);
		verts.insert(lca(gang[c][i], gang[c][i+1]));
	}
	vector<int> mvt;
	for(auto u : verts) {
		mvt.push_back(u);
		virtuals[u].push_back(c);
		v_childs[u].push_back({});
	}
	sort(mvt.begin(), mvt.end(), [&](int a, int b){
		return pre[a] < pre[b];
	});
	vector<int> stos;
	for(auto v : mvt) {
		while(stos.size() && lca(stos.back(), v)!=stos.back()) stos.pop_back();
		if(stos.size()) v_childs[stos.back()].back().push_back(v);
		stos.push_back(v);
	}
	grab_ans[mvt[0]].push_back(c);
}

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

	cin >> n >> m;
	for(int i=1; i<n; ++i) {
		int a, b;
		cin >> a >> b;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	for(int i=1; i<=n; ++i) cin >> C[i];
	path_end[1] = 1;
	jump_dfs(1,0);
	make_hld(1,0);
	for(int i=1; i<=n; ++i) gang[C[i]].push_back(i);
	for(int i=1; i<=m; ++i) {
		make_vt(i);
	}
	mw_dfs(1, 0);
	for(int i=1; i<=m; ++i) cout << ans[i] << " "; cout << "\n"; 
	int res = n;
	for(int i=1; i<=m; ++i) {
		int ile = 0;
		while(ans[i] > 1) {
			ile += ans[i]/2;
			ans[i] = (ans[i] + 1)/2;
		}
		res = min(res, ile);
	}
	cout << res << "\n";

	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'ds mw_dfs(int, int)':
capital_city.cpp:131:39: error: 'u' was not declared in this scope
  131 |                         if(d[path_end[u]] <= d[v]) {
      |                                       ^
capital_city.cpp:138:42: error: 'u' was not declared in this scope
  138 |                         if(v==2) cout << u << ", ";
      |                                          ^