Submission #393046

#	Time	Username	Problem	Language	Result	Execution time	Memory
393046		Tony1234	Construction of Highway (JOI18_construction)	C++14	0 / 100	53 ms	70416 KiB

construction

#include<bits/stdc++.h>
using namespace std;
typedef pair<int, int> pii;
const int mx = 1e5+10;
int dep[mx], chainid[mx], chainind[mx], par[mx], sz[mx], head[mx];
vector<int> adj[mx];
int val[mx];
deque<pii> chains[mx];
void dfs(int cur, int pa){
	par[cur] = pa;
	sz[cur] = 1;
	dep[cur] = (pa == -1 ? 0 : dep[pa] + 1);
	for(int u: adj[cur]){
		if(u^pa){
			dfs(u, cur);
			sz[cur] += sz[u];	
		}
	}
}
int numchains = 1;
void hld(int cur, int pa, int idx){
	if(head[numchains] == -1){
		head[numchains] = cur;
	}
	chainid[cur] = numchains;
	if(chains[numchains].empty()){
		chains[numchains].push_back({val[cur], 1});
	}else{
		if(chains[numchains].back().first == val[cur])chains[numchains].back().second++;
		else chains[numchains].push_back({val[cur], 1});
	}
	//cout << "ijewafaijefojaweoijfiw\n";
	int spec = -1;
	for(int u: adj[cur]){
		if(u^pa){
			if(spec == -1 || sz[u] >= sz[spec]){
				spec= u;
			}
		}
	}
	if(spec != -1)hld(spec, cur, idx+1);
	for(int u: adj[cur]){
		if( u ^ pa && u ^ spec){
			numchains++;
			hld(u, cur, idx);
		}
	}
}
int bit[mx];
void upd(int idx, int val){
	for(; idx < mx; idx += idx & -idx){
		bit[idx] += val;
	}
}
int query(int idx){
	int ans = 0;
	for(; idx; idx -= idx & -idx){
		ans += bit[idx];
	}
	return ans;
}
int query_path(int u){
	//return 69;
	
	vector<pii> qq; vector<int> ss; //need to resize it later
	while(u){
		int nxt = head[chainid[u]];
		int sz = dep[u] - dep[nxt] + 1;
		vector<pii> add;
	//	cout << nxt << " " << sz << "\n";
		for(auto u: chains[chainid[nxt]]){
			if(u.second >= sz){
				add.push_back({u.first, sz});
				break;
			}else{
				add.push_back(u);
				sz -= u.second;
			}
		}
		reverse(add.begin(), add.end());
		for(auto cc: add){
			ss.push_back(cc.first);
			qq.push_back(cc);
		}
		u = par[nxt];
	}
	reverse(qq.begin(), qq.end());
	ss.erase(unique(ss.begin(), ss.end()), ss.end());
	sort(ss.begin(), ss.end());
/*	for(auto u1 : ss){
		cout << u1 << "\n";
	}
	for(auto u1: qq){
		cout << u1.first << " " << u1.second << "\n"; 
	}*/
	long long ans = 0, siz = 0;
	for(int i=1; i<= ss.size() + 100; i++){
		bit[i] = 0;
	}
	for(pii cc : qq){
		int ind = lower_bound(ss.begin(), ss.end(), cc.first) - ss.begin() + 1;
		ans += 1LL * (siz - query(ind)) * cc.second;
		siz += cc.second;
		upd(ind, cc.second);
	}
	return ans;
}
void link_path(int u, int newval){
	//update all the chains from u --> root, for each deque, we can keep on popping the front element
	while(u){
		int headchain = head[chainid[u]]; //this is the head of the chain
		int siz = dep[u] - dep[headchain] + 1;
	//	cout << siz << " lol\n";
		while(siz && chains[chainid[u]].size()){
			if(chains[chainid[u]].front().second > siz){
				chains[chainid[u]].front().second -= siz;
				siz = 0;
			}else{
				siz -= chains[chainid[u]].front().second;
				chains[chainid[u]].pop_front();
			}
		}
		//now we need to push the front
		chains[chainid[u]].push_front({newval, dep[u] - dep[headchain] + 1});
		u = par[head[chainid[u]]];
	}
}
void debug(){
	cout << numchains << "\n";
	for(int i=1; i <= numchains; i++){
		cout << i << "\n";
		for(auto cc : chains[i]){
			cout << cc.first << " " << cc.second << "\n";
		}
	}
}
int main(){
	ios_base::sync_with_stdio(false); cin.tie(0);
	memset(head, -1, sizeof(head));
	int n; cin >> n;
	for(int i=1; i <= n; i++)cin >> val[i];
	vector<pair<int, int>> edges;
	for(int i=0; i < n-1; i++){
		int a, b; cin >> a >> b; adj[a].push_back(b); adj[b].push_back(a); edges.push_back(make_pair(a, b));
		//cout << a << " " << b << endl;
	}
	

	dfs(1, -1);
	//cout << "Finished dfs\n";

	hld(1, -1, -1);
//		debug();	

	for(auto u: edges){
//		debug();
		cout << query_path(u.first) << "\n";
		link_path(u.first, val[u.second]);
	}
}

Compilation message (stderr)

construction.cpp: In function 'int query_path(int)':
construction.cpp:97:16: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   97 |  for(int i=1; i<= ss.size() + 100; i++){
      |               ~^~~~~~~~~~~~~~~~~~

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