Submission #1227543

#	Time	Username	Problem	Language	Result	Execution time	Memory
1227543		_rain_	Construction of Highway (JOI18_construction)	C++20	100 / 100	241 ms	108224 KiB

construction

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

const int N=(int)1e5;
const int MAXLOG=17;
	vector<int>ke[N+2];
	vector<int>nen;
	void add_canh(int u,int v){
		ke[u].push_back(v) , ke[v].push_back(u);
		return;
	}
	int u[N+2],v[N+2],c[N+2];
	int n;
	
	class Fenwick{
		public:
			vector<int>bit;
			int n;
			void init(int _n){
				n=_n;
				bit.assign(n+2,0);
				return;
			}
			void upd(int pos,int val){
				for(;pos<=n;pos+=pos&-pos) bit[pos]+=val;
				return;
			}
			int Get(int pos){
				int sum=0; 
				for(;pos;pos-=pos&-pos) sum+=bit[pos];
				return sum;
			}
	};
	Fenwick tree;
	
namespace HLD{
	
	int son[N+2],sz[N+2],par[N+2];
	int chainID[N+2],head[N+2],dep[N+2]={};
	
	int cur_chain=1;
	
	void pre_build(int u,int p){
		par[u]=p;
		sz[u]=1;
		for(auto&v:ke[u]){
			if (v==p) continue;
			pre_build(v,u);
			sz[u]+=sz[v];
		}
		for(auto&v:ke[u]){
			if (v==p) continue;
			if (sz[son[u]]<sz[v]) son[u]=v;
		}
	}
	void build(int u,int p){
		chainID[u]=cur_chain;
		dep[u]=dep[p]+1;
		if (head[cur_chain]==0) head[cur_chain]=u;
		if (son[u]) build(son[u],u);
		for(auto&v:ke[u]){
			if (v==p||v==son[u]) continue;
			++cur_chain;
			build(v,u);
		}
	}
	
	deque<pair<int,int>>bucket[N+2];	
	
	void insert(int a,int b){
		int color=c[b],u=b;
		while (chainID[u]){
			int need_h=dep[u]-dep[head[chainID[u]]]+1;
			int sz=need_h;
			while (bucket[chainID[u]].size() && need_h){
				auto &[x1,x2]=bucket[chainID[u]].front();
				if (x2<=need_h){
					need_h-=x2;
					bucket[chainID[u]].pop_front();
				} 
				else {
					x2-=need_h;
					break;
				}
			}
			bucket[chainID[u]].push_front({color,sz});
			u=par[head[chainID[u]]];
		}
		return;
	}
	
	LL calc(int a){
		int u=a;
		vector<pair<int,int>>v;
		vector<int>color;
		while (chainID[u]>0){
			int need_h=dep[u]-dep[head[chainID[u]]]+1;
			vector<pair<int,int>>tst;
			for(int i=0;i<bucket[chainID[u]].size() && need_h;++i){
				auto [x1,x2]=bucket[chainID[u]][i];
				color.push_back(x1);
				if (need_h>=x2){
					tst.push_back(bucket[chainID[u]][i]);
					need_h-=x2;
				}
				else {
					tst.push_back({x1,need_h});
					need_h=0;
				}
			}
			for(int i=tst.size()-1;i>=0;--i) v.push_back(tst[i]);
			u=par[head[chainID[u]]];
		}
		sort(color.begin(),color.end());
		color.resize(unique(color.begin(),color.end())-color.begin());
		reverse(v.begin(),v.end());
		tree.init(color.size());
		LL total=0 , ans=0;
		for(auto&x:v){
			int toado=upper_bound(color.begin(),color.end(),x.first)-color.begin();
			ans+=(LL)x.second*(total-tree.Get(toado));
			total+=x.second;
			tree.upd(toado,x.second);
		}
		return ans;
	}
};
using namespace HLD;

int main(){
	ios::sync_with_stdio(false);
	cin.tie(0) ; cout.tie(0);
	#define task "main"
	if (fopen(task".inp","r")){
		freopen(task".inp","r",stdin);
		freopen(task".out","w",stdout);
	}

	cin>>n;
	for(int i=1;i<=n;++i) cin>>c[i];
	for(int i=1;i<n;++i){
		cin>>u[i]>>v[i];
		add_canh(u[i],v[i]);
	}
	pre_build(1,0) , build(1,0);
	for(int i=1;i<n;++i){
		cout<<calc(u[i])<<'\n';
		insert(u[i],v[i]);
	}
	return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:136:24: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  136 |                 freopen(task".inp","r",stdin);
      |                 ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~
construction.cpp:137:24: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  137 |                 freopen(task".out","w",stdout);
      |                 ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~

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