Submission #595660

#	Time	Username	Problem	Language	Result	Execution time	Memory
595660		Bobaa	Construction of Highway (JOI18_construction)	C++17	100 / 100	443 ms	93876 KiB

construction

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

const int maxn = 1e5 + 5; 

int lv[maxn], sz[maxn], par[maxn], tree[maxn], in[maxn], out[maxn], tin, sizeofchain[maxn]; 
map<int, int> mp; 
vector<int> v[maxn], tocmp; 
deque<pair<int, int>> vec[maxn];
int bigchild[maxn], chain[maxn], n, a[maxn], b[maxn], c[maxn]; 

void upd(int idx, int val) {
	for (int i = idx; i < maxn; i+= i & (-i)) tree[i] += val; 
} 

int solve(int idx) {
	int pas = 0; 
	for (int i = idx; i > 0; i -= i & (-i)) pas += tree[i]; 
	return pas;
}

void dfs1(int a, int p) {
	sz[a] = 1; 
	lv[a] = lv[p] + 1; 
	par[a] = p; 
	for (int i = 0; i < v[a].size(); i++) {
		int to = v[a][i]; 
		if (to == p) continue; 
		dfs1(to, a); 
		sz[a] += sz[to]; 
		if (!bigchild[a] || sz[bigchild[a]] < sz[to]) bigchild[a] = to; 
	}
}

void dfs2(int a, int p) {
	tin++; 
	in[a] = tin; 
	if (bigchild[a]) dfs2(bigchild[a], a); 
	for (int i = 0; i < v[a].size(); i++) {
		int to = v[a][i]; 
		if (to == p || to == bigchild[a]) continue; 
		dfs2(to, a); 
	}
	out[a] = tin; 
}

void chain_dfs(int a, int p) {
	if (bigchild[a]) chain[bigchild[a]] = chain[a]; 
	sizeofchain[a]++; 
	if (!vec[chain[a]].size()) vec[chain[a]].push_back({c[a], 1}); 
	else {
		if (vec[chain[a]].back().first == c[a]) vec[chain[a]].back().second++; 
		else vec[chain[a]].push_back({c[a], 1}); 
	}
	for (int i = 0; i < v[a].size(); i++) {
		int to = v[a][i]; 
		if (to == p) continue; 
		chain_dfs(to, a); 
	}
}

int solve(int a, int b) {
	int res = 0; 
	int lst = chain[a]; 
	int cur = a; 
	vector<pair<int, int>> pas; 
	while (true) {
		lst = chain[cur]; 
		int x = lv[cur] - lv[lst] + 1; 
		vector<pair<int, int>> vv; 
		int sum = 0; 
		for (int i = 0; i < vec[lst].size(); i++) {
			if (vec[lst][i].second + sum > x) {
				vv.push_back({vec[lst][i].first, (x - sum)}); 
				break; 
			}
			sum += vec[lst][i].second; 
			vv.push_back({vec[lst][i].first, vec[lst][i].second}); 
			if (sum == x) break; 
		}
		reverse(vv.begin(), vv.end()); 
		for (int i = 0; i < vv.size(); i++) pas.push_back({vv[i].first, vv[i].second}); 
		cur = par[lst]; 
		if (cur == 0) break; 
	}
	for (int i = 0; i < pas.size(); i++) {
		res += solve(pas[i].first - 1) * pas[i].second; 
		upd(pas[i].first, pas[i].second); 
	}
	for (int i = 0; i < pas.size(); i++) upd(pas[i].first, -pas[i].second); 
	return res; 
}

void toadd(int a, int b) {
	int lst = chain[a]; 
	int cur = a;
	while (true) {
		lst = chain[cur]; 
		int x = lv[cur] - lv[lst] + 1; 
		int sum = 0; 
		vector<pair<int, int>> vv; 
		while (vec[lst].size()) {
			int i = 0; 
			if (vec[lst][i].second + sum > x) {
				vec[lst][i].second = vec[lst][i].second - (x - sum); 
				break; 
			}
			sum += vec[lst][i].second; 
			vec[lst].pop_front(); 
			if (sum == x) break; 
		}
		vec[lst].push_front({c[b], x}); 
		cur = par[lst]; 
		if (cur == 0) break; 
	}
}

int main() {
	ios_base::sync_with_stdio(false); 
	cin.tie(nullptr); 
	cin >> n; 
	for (int i = 1; i <= n; i++) {
		cin >> c[i]; 
		tocmp.push_back(c[i]); 
	}
	sort(tocmp.begin(), tocmp.end()); 
	int cnt = 1; 
	mp[tocmp[0]] = 1; 
	for (int i = 1; i < tocmp.size(); i++) {
		if (tocmp[i] != tocmp[i - 1]) cnt++; 
		mp[tocmp[i]] = cnt; 
	}
	for (int i = 1; i <= n; i++) c[i] = mp[c[i]]; 
	for (int i = 1; i <= n - 1; i++) {
		cin >> a[i] >> b[i]; 
		v[a[i]].push_back(b[i]); 
		v[b[i]].push_back(a[i]); 
	}
	for (int i = 1; i <= n; i++) chain[i] = i; 
	dfs1(1, 0); 
	dfs2(1, 0); 
	chain_dfs(1, 0); 
	for (int i = 1; i <= n - 1; i++) {
		int ans = solve(a[i], b[i]); 
		toadd(a[i], b[i]); 
		cout << ans << '\n'; 
	}
}

Compilation message (stderr)

construction.cpp: In function 'void dfs1(int, int)':
construction.cpp:26:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   26 |  for (int i = 0; i < v[a].size(); i++) {
      |                  ~~^~~~~~~~~~~~~
construction.cpp: In function 'void dfs2(int, int)':
construction.cpp:39:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   39 |  for (int i = 0; i < v[a].size(); i++) {
      |                  ~~^~~~~~~~~~~~~
construction.cpp: In function 'void chain_dfs(int, int)':
construction.cpp:55:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   55 |  for (int i = 0; i < v[a].size(); i++) {
      |                  ~~^~~~~~~~~~~~~
construction.cpp: In function 'int solve(int, int)':
construction.cpp:72:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::deque<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   72 |   for (int i = 0; i < vec[lst].size(); i++) {
      |                   ~~^~~~~~~~~~~~~~~~~
construction.cpp:82:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   82 |   for (int i = 0; i < vv.size(); i++) pas.push_back({vv[i].first, vv[i].second});
      |                   ~~^~~~~~~~~~~
construction.cpp:86:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   86 |  for (int i = 0; i < pas.size(); i++) {
      |                  ~~^~~~~~~~~~~~
construction.cpp:90:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   90 |  for (int i = 0; i < pas.size(); i++) upd(pas[i].first, -pas[i].second);
      |                  ~~^~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:129:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  129 |  for (int i = 1; i < tocmp.size(); i++) {
      |                  ~~^~~~~~~~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84