This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
using namespace std;
typedef tuple<int,int,int> trinca;
typedef pair<int,int> ii;
const int MAXN = 1e5 + 10;
vector<int> grafo[MAXN];
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;
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()){
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++){
for(int j = i+1; j < pares.size(); j++){
if(pares[i].first > pares[j].first){
ans += 1LL*pares[i].second*pares[j].second;
}
}
}
return ans;
}
int main(){
scanf("%d", &N);
for(int i = 1; i <= N; i++){
scanf("%d", &color[i]);
}
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:119:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i = 0; i < pares.size(); i++){
~~^~~~~~~~~~~~~~
construction.cpp:120:22: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int j = i+1; j < pares.size(); j++){
~~^~~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:133:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
scanf("%d", &N);
~~~~~^~~~~~~~~~
construction.cpp:136:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
scanf("%d", &color[i]);
~~~~~^~~~~~~~~~~~~~~~~
construction.cpp:140: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]);
~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |