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;
#define int long long
#define pii pair<int, int>
vector<int> colors;
vector<int> time_g;
vector<int> depth;
map<int, int> m;
struct Node{
int t;
int last_m;
int color;
mutable vector<pii> cols;
bool operator<( const Node& b)const{
return t<b.t;
}
Node(int id){
t = time_g[id];
color = colors[id];
last_m =id;
//cols.push_back(pii(depth[id]-1, color));
}
};
vector<vector<pii>> adj;
vector<set<Node>> sets;
void merge_pure(set<Node>&s, set<Node>&b, int prov){
for(auto e: b){
auto e2 = e;
e2.last_m = prov;
s.insert(e);
}
}
void compute_rel(set<Node>::iterator a, set<Node>::iterator b, int d){
//cout<<a->t<<" "<<b->t<<endl;
if(a->last_m!=b->last_m){
//cout<<m[a->t]+1<<" includes "<<m[b->t]+1<<endl;
a->cols.push_back(pii(d, b->color));
}
}
void merge(set<Node>&a, set<Node>& b, int id){
for(auto it = b.begin(); it!=b.end(); it++){
auto pre = a.lower_bound((*it));
if(pre!=a.end()){
it++;
if(it== b.end() || pre->t<it->t){
it--;
compute_rel(pre, it, depth[id]);
it++;
}
it--;
}
if(pre!=a.begin()){
pre--;
if(pre!=a.end() && ((*pre)<(*it))){
compute_rel(it, pre, depth[id]);
}
}
a.insert(*it);
}
}
void dfs(int id){
set<Node> res;
set<Node> comp;
if(adj[id].size()>0){
pii best = pii(0, 0);
for(auto ch: adj[id]){
dfs(ch.first);
if(sets[ch.first].size()>best.first){
best.first = sets[ch.first].size();
best.second = ch.first;
}
}
swap(res, sets[best.second]);
for(auto e: adj[id]){
if(e.first != best.second){
merge_pure(comp, sets[e.first],e.first);
}
}
}
comp.insert(Node(id));
merge(res, comp, id);
swap(res, sets[id]);
}
int merge(vector<pii>&v, int lt, int rt, int mid){
vector<pii> res;
int c= 0;
int l = lt;
int r= mid+1;
int r_s =0;
for(int i = lt; i<=rt; i++){
if(l>mid){
res.push_back(v[r]);
r_s += v[r].first;
r++;
}
else if(r>rt){
res.push_back(v[l]);
c+=(r_s)*v[l].first;
l++;
}
else if(v[l].second<=v[r].second){
res.push_back(v[l]);
c+=(r_s)*v[l].first;
l++;
}
else{
res.push_back(v[r]);
r_s += v[r].first;
r++;
}
}
for(int i = lt; i<=rt; i++){
v[i] = res[i-lt];
}
return c;
}
int dpr(vector<pii>&v, int lt, int rt){
if(lt==rt){
return 0;
}
else{
//cout<<lt<<" "<<rt<<endl;
int mid = (lt+rt)/2;
int r =0;
r+=dpr(v, lt, mid);
r+=dpr(v, mid+1, rt);
r+= merge(v, lt, rt, mid);
//cout<<r<<endl;
return r;
}
}
signed main(){
int n;
cin>>n;
colors.resize(n);
adj.resize(n);
time_g.resize(n, -1);
depth.resize(n, 1);
sets.resize(n);
for(int i = 0; i<n; i++){
int c;
cin>>c;
colors[i] = c;
}
m.insert(pii(-1, 0));
for(int i = 0; i<n-1; i++){
int a, b;
cin>>a>>b;
a--;b--;
time_g[b] = i;
adj[a].push_back(pii(b, i));
depth[b] = depth[a]+1;
m.insert(pii(i, b));
}
dfs(0);
set<Node> res;
swap(res, sets[0]);
vector<int> ans(n-1);
for(auto e: res){
int cur= e.t;
if(cur!=-1){
int r= 0;
sort(e.cols.begin(), e.cols.end());
for(int i = e.cols.size()-1; i>0; i--){
e.cols[i].first -=e.cols[i-1].first;
}
r += dpr(e.cols, 0, e.cols.size()-1);
ans[cur]= r;
}
}
for(int i= 0; i<n-1; i++){
cout<<ans[i]<<endl;
}
}
Compilation message (stderr)
construction.cpp: In function 'void dfs(long long int)':
construction.cpp:79:37: warning: comparison of integer expressions of different signedness: 'std::set<Node>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
79 | if(sets[ch.first].size()>best.first){
| ~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |