Submission #256425

#	Time	Username	Problem	Language	Result	Execution time	Memory
256425		IvanC	Construction of Highway (JOI18_construction)	C++17	100 / 100	406 ms	85868 KiB

construction

#include <bits/stdc++.h>
#define LSOne(S) (S&(-S))
using namespace std;

typedef tuple<int,int,int> trinca;
typedef pair<int,int> ii;

const int MAXN = 1e5 + 10;

vector<int> grafo[MAXN];
vector<int> compressao;
vector<ii> quantities;
deque<trinca> componente[MAXN];
int pai[MAXN], szTree[MAXN], color[MAXN], aresta1[MAXN], aresta2[MAXN], N;
int chainId[MAXN], posChain[MAXN], headChain[MAXN], chainPtr, lastChain;
int fenwickTree[MAXN];

inline void upd_bit(int pos, int delta){
	while(pos < MAXN){
		fenwickTree[pos] += delta;
		pos += LSOne(pos);
	}
}

inline int read(int pos){
	int ans = 0;
	while(pos > 0){
		ans += fenwickTree[pos];
		pos -= LSOne(pos);
	}
	return ans;
}

int query(int pos){
	return read(MAXN-1) - read(pos);
}

void dfs1(int v, int p){

	szTree[v] = 1;
	pai[v] = p;

	for(int u: grafo[v]){
		if(u == p) continue;
		dfs1(u, v);
		szTree[v] += szTree[u];
	}

}

void dfs2(int v, int p){

	int mx = -1, big = -1;
	if(chainId[v] == 0){
		chainId[v] = ++lastChain;
		headChain[chainId[v]] = v;
		posChain[v] = ++chainPtr;
	}

	componente[chainId[v]].push_back(make_tuple(color[v], posChain[v], posChain[v]));

	for(int u: grafo[v]){
		if(u == p) continue;
		if(szTree[u] > mx){
			mx = szTree[u];
			big = u;
		}
	}

	if(big != -1){
		chainId[big] = chainId[v];
		posChain[big] = ++chainPtr;
		dfs2(big, v);
	}

	for(int u: grafo[v]){
		if(u == p || u == big) continue;
		dfs2(u, v);
	}

}

long long query_up(int v){

	//printf("#######\nQuery %d\n", v);
	vector<ii> pares;
	int new_color = color[v];
	v = pai[v];
	long long ans = 0;

	while(v != -1){

		int thisChain = chainId[v];
		int thisHead = headChain[thisChain];
		int lo = posChain[thisHead];
		int hi = posChain[v];

		//printf("V %d tC %d tH %d (%d, %d)\n", v, thisChain, thisHead, lo, hi);

		while(!componente[thisChain].empty()){
			//printf("Loop\n");
			trinca davez = componente[thisChain].front();
			if(get<1>(davez) > hi){
				break;
			}
			else if(get<2>(davez) <= hi){
				int whichColor = get<0>(davez);
				int whichSize = get<2>(davez) - get<1>(davez) + 1;
				componente[thisChain].pop_front();
				quantities.push_back(ii(whichColor, whichSize));
			}
			else{
				int whichColor = get<0>(davez);
				int oldBegin = get<1>(davez);
				int oldEnd = get<2>(davez);
				componente[thisChain].pop_front();
				componente[thisChain].push_front(make_tuple(whichColor, hi + 1, oldEnd));
				int whichSize = hi - oldBegin + 1;
				quantities.push_back(ii(whichColor, whichSize));
			}
		}

		while(!quantities.empty()){
			if(!pares.empty() && pares.back().first == quantities.back().first){
				pares.back().second += quantities.back().second;
			}
			else{
				pares.push_back(quantities.back());
			}
			quantities.pop_back();
		}

		componente[thisChain].push_front(make_tuple(new_color, lo, hi));

		v = pai[thisHead];

	} 

	reverse(pares.begin(), pares.end());
	/*
	for(int i = 0; i < pares.size(); i++){
		printf("(%d, %d)", pares[i].first, pares[i].second);
	}
	printf("\n");
	printf("#######\n");*/

	for(int i = 0; i < pares.size(); i++){
		int thisColor = pares[i].first;
		int thisQuantity = pares[i].second;
		ans += 1LL*thisQuantity*query(thisColor);
		upd_bit(thisColor, thisQuantity);
	}

	for(int i = 0; i < pares.size(); i++){
		int thisColor = pares[i].first;
		int thisQuantity = pares[i].second;
		upd_bit(thisColor, -thisQuantity);
	}

	return ans;

}

int main(){

	scanf("%d", &N);

	for(int i = 1; i <= N; i++){
		scanf("%d", &color[i]);
		compressao.push_back(color[i]);
	}
	sort(compressao.begin(), compressao.end());
	for(int i = 1; i <= N; i++){
		color[i] = lower_bound(compressao.begin(), compressao.end(), color[i]) - compressao.begin() + 1;
	}

	for(int i = 1; i < N; i++){
		scanf("%d %d", &aresta1[i], &aresta2[i]);
		grafo[aresta1[i]].push_back(aresta2[i]);
	}

	dfs1(1, -1);
	dfs2(1, -1);

	for(int i = 1; i < N; i++){

		int v = aresta2[i];
		long long ans = query_up(v);
		printf("%lld\n", ans);

	}

	return 0;

}

Compilation message (stderr)

construction.cpp: In function 'long long int query_up(int)':
construction.cpp:147:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for(int i = 0; i < pares.size(); i++){
                 ~~^~~~~~~~~~~~~~
construction.cpp:154:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for(int i = 0; i < pares.size(); i++){
                 ~~^~~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:166:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d", &N);
  ~~~~~^~~~~~~~~~
construction.cpp:169:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d", &color[i]);
   ~~~~~^~~~~~~~~~~~~~~~~
construction.cpp:178:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d", &aresta1[i], &aresta2[i]);
   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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