Submission #412741

#	Time	Username	Problem	Language	Result	Execution time	Memory
412741		errorgorn	Road Closures (APIO21_roads)	C++17	100 / 100	517 ms	344608 KiB

roads

#include "roads.h"

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

#define ll long long
#define ii pair<ll,ll>
#define fi first
#define se second

#define puf push_front
#define pof pop_front
#define pub push_back
#define pob pop_back

#define rep(x,s,e) for (auto x=s-(s>e);x!=e-(s>e);(s<e?x++:x--))
#define all(x) (x).begin(),(x).end()
#define sz(x) (int) (x).size()

struct node{
	int s,e,m;
	ll cnt=0,val=0;
	node *l=nullptr,*r=nullptr;
	
	node (int _s,int _e){
		s=_s,e=_e,m=s+e>>1;
	}
	
	void update(int i,int j,int k){
		cnt+=j,val+=k;
		
		if (s==e) return;
		if (i<=m){
			if (l==nullptr) l=new node(s,m);
			l->update(i,j,k);
		}
		else{
			if (r==nullptr) r=new node(m+1,e);
			r->update(i,j,k);
		}
	}
	
	ll query(ll num){
		if (s==e) return num*s;
		if (cnt==num) return val;
		if (l==nullptr) return r->query(num);
		else if (l->cnt>=num) return l->query(num);
		else return l->val+r->query(num-l->cnt);
	}
} *root[100005];

ll n;
set<ii> al[100005];

ll k;

vector<ll> proc;
ll ss[100005];
bool can[100005];

ll vis[100005];

//how many dead people have a certain weight

ii dfs(ll i,ll p){
	vis[i]=k;
	vector<ll> deltas;
	ll tot=0;
	ll num=0;
	
	for (auto &it:al[i]){
		if (it.fi==p) continue;
		if (!can[it.fi]) continue;
		
		ii temp=dfs(it.fi,i);
		
		//cout<<"dfs: "<<i<<" "<<it.fi<<" "<<temp.fi<<" "<<temp.se<<endl;
		
		tot+=temp.se;
		ll delta=temp.fi-temp.se+it.se;
		if (delta<0) tot+=delta,num++;
		else deltas.pub(delta);
	}
	
	ii ret;
	
	for (auto &it:deltas) root[i]->update(it,1,it);
	
	//cout<<i<<" "<<p<<" "<<root[i]->cnt<<" "<<ss[i]-k<<" "<<num<<endl;
	if (ss[i]-k-1-num<=0) ret.fi=tot;
	else ret.fi=tot+root[i]->query(ss[i]-k-1-num);
	
	if (ss[i]-k-num<=0) ret.se=tot;
	else ret.se=tot+root[i]->query(ss[i]-k-num);
	
	for (auto &it:deltas) root[i]->update(it,-1,-it);
	
	return ret;
}

std::vector<long long> minimum_closure_costs(int N, std::vector<int> U,
                                             std::vector<int> V,
                                             std::vector<int> W) {
				
	n=N;
	ll tot=0;
	rep(x,0,n-1){
		al[U[x]].insert(ii(V[x],W[x]));
		al[V[x]].insert(ii(U[x],W[x]));
		tot+=W[x];
	}
	
	memset(vis,-1,sizeof(vis));
						 
	vector<ll> ans(n,0);
	ans[0]=tot;
	
	rep(x,0,n) ss[x]=sz(al[x]);
	rep(x,0,n) proc.pub(x),can[x]=true;
	sort(all(proc),[](int i,int j){
		return ss[i]>ss[j];
	});
	
	rep(x,0,n) root[x]=new node(0,1e9+5);
	
	rep(x,1,n){
		k=x;
		
		while (!proc.empty() && ss[proc.back()]<=x){
			int node=proc.back();
			//cout<<"delete: "<<node<<endl;
			can[node]=false;
			
			for (auto &it:al[node]){
				al[it.fi].erase(ii(node,it.se));
				root[it.fi]->update(it.se,1,it.se);
			}
			
			proc.pob();
		}
		
		ll val=0;
		for (auto &it:proc){
			//cout<<"processing: "<<x<<" "<<it<<endl;
			if (vis[it]!=k) val+=dfs(it,-1).se;
		}
		
		ans[x]=val;
	}
	
	return ans;
}

Compilation message (stderr)

roads.cpp: In constructor 'node::node(int, int)':
roads.cpp:26:16: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   26 |   s=_s,e=_e,m=s+e>>1;
      |               ~^~

• Subtask 1: 5 / 5
• Subtask 2: 7 / 7
• Subtask 3: 14 / 14
• Subtask 4: 10 / 10
• Subtask 5: 17 / 17
• Subtask 6: 25 / 25
• Subtask 7: 22 / 22