Submission #706702

#TimeUsernameProblemLanguageResultExecution timeMemory
706702attemptn63Construction of Highway (JOI18_construction)C++17
100 / 100
336 ms112900 KiB
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define int long long
#define double long double
#define endl '\n'
#define all(k) (k).begin(), (k).end()
#define pb push_back
#define fi first
#define se second
#define sz(k) (int)(k).size()
#define speed ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0)
#define pi pair<long long, long long>
#define mp make_pair
#define ti tuple<long long, long long, long long>
#define INF (long long)1e18
#define chmin(a, b) a = min(a, b)
#define chmax(a, b) a = max(a, b)
#define sign(x) ((x > 0) - (x < 0))
template <class T>using vec = vector<T>;
template <class T>using min_heap = priority_queue<T, vec<T>, greater<T>>;
template <class T>using max_heap = priority_queue<T>;
template <class T>using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
template <class T>using ordered_multiset = tree<T, null_type, less_equal<T>, rb_tree_tag, tree_order_statistics_node_update>;
#ifdef DEBUG
template <typename T>
inline void debug(T k){
	cout << k << " ";
}
template <typename T, typename... Args>
inline void debug(T k, Args... args){
	cout << k << " | ";
	debug(args...);
}
#define debug(...) cout << "[" << #__VA_ARGS__ << "] : ";debug(__VA_ARGS__);cout << endl
template <typename T>
ostream &operator<<(ostream &o_str, const vec<T> &p){
	o_str << "{ ";
	for (auto i = 0; i < (int)p.size(); i++){
		o_str << p[i];
		if (i < (int)p.size() - 1){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T>
ostream &operator<<(ostream &o_str, const deque<T> &p){
	o_str << "{ ";
	for (auto i = 0; i < (int)p.size(); i++){
		o_str << p[i];
		if (i < (int)p.size() - 1){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const pair<T, U> &p){
	return o_str << "(" << p.fi << ", " << p.se << ")";
}
template <typename T, typename U, typename V>
ostream &operator<<(ostream &o_str, const tuple<T, U, V> &p){
	return o_str << "(" << get<0>(p) << ", " << get<1>(p) << ", " << get<2>(p) << ")";
}
template <typename T>
ostream &operator<<(ostream &o_str, const set<T> &p){
	o_str << "{ ";
	for (auto i = p.begin(); i != p.end(); i++){
		o_str << *i;
		if (i != --p.end()){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const map<T, U> &p){
	o_str << "{ ";
	for (auto i = p.begin(); i != p.end(); i++){
		o_str << i->first << ": " << i->second;
		if (i != --p.end()){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T>
ostream &operator<<(ostream &o_str, const queue<T> &p){
	queue<T> q = p;
	o_str << "{ ";
	while (!q.empty()){
		o_str << q.front();
		q.pop();
		if (!q.empty()){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T, typename U, typename cmp>
ostream &operator<<(ostream &o_str, const priority_queue<T,U,cmp> &p){
	priority_queue<T,U,cmp> q = p;
	o_str << "{ ";
	while (!q.empty()){
		o_str << q.top();
		q.pop();
		if (!q.empty()){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const set<T, U> &p){
	o_str << "{ ";
	for (auto i = p.begin(); i != p.end(); i++){
		o_str << *i;
		if (i != --p.end()){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template<typename T, typename U>
ostream &operator<<(ostream &o_str, const multiset<T, U> &p){
	o_str << "{ ";
	for (auto i = p.begin(); i != p.end(); i++){
		o_str << *i;
		if (i != --p.end()){
			o_str << ", ";
		}
	}
	return o_str << " }" << endl;
}
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const unordered_map<T, U> &p){
	o_str << "{ ";
	for (auto i = p.begin(); i != p.end(); i++){
		o_str << i->first << ": " << i->second;
		o_str << ", ";
	}
	return o_str << " }" << endl;
}
#else
template <typename T>
inline void debug(T k){}
template <typename T, typename... Args>
inline void debug(T k, Args... args) {}
#define debug(...)
template <typename T>
ostream &operator<<(ostream &o_str, const vec<T> &p){return o_str;}
template <typename T>
ostream &operator<<(ostream &o_str, const deque<T> &p) { return o_str; }
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const pair<T, U> &p) { return o_str; }
template <typename T, typename U, typename V>
ostream &operator<<(ostream &o_str, const tuple<T, U, V> &p) { return o_str; }
template <typename T>
ostream &operator<<(ostream &o_str, const set<T> &p) { return o_str; }
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const map<T, U> &p) { return o_str; }
template <typename T>
ostream &operator<<(ostream &o_str, const queue<T> &p) { return o_str; }
template <typename T, typename U, typename cmp>
ostream &operator<<(ostream &o_str, const priority_queue<T,U,cmp> &p){return o_str;}
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const set<T, U> &p) { return o_str; }
template<typename T, typename U>
ostream &operator<<(ostream &o_str, const multiset<T, U> &p) { return o_str; }
template <typename T, typename U>
ostream &operator<<(ostream &o_str, const unordered_map<T, U> &p) { return o_str; }
#endif

int n, val[100005],depth[100005],sz[100005],hpar[100005],par[100005];
deque<pi> hld[100005];
vec<int> adj[100005],disc;
vec<pi> edges,chain;
void dfs_sz(int i){
	sz[i] = 1;
	for(auto &j: adj[i]){
		par[j] = i;
		depth[j] = depth[i] + 1;
		dfs_sz(j);
		sz[i] += sz[j];
		if(sz[j] > sz[adj[i][0]]) swap(j,adj[i][0]);
	}
}
void dfs_hld(int i){
	for(auto it: adj[i]){
		if(it == adj[i][0]) hpar[it] = hpar[i];
		else hpar[it] = it;
		dfs_hld(it);
	}
}

int bit[100005];
void init(int n){
	for(int i = 0;i<=n;i++) bit[i] = 0;
}
void upd(int i,int v){
	i++;
	for(;i<=n;i+=(i&-i)) bit[i] += v;
}
int qry(int i){
	i++;
	int ret = 0;
	for(;i;i-=(i&-i)) ret += bit[i];
	return ret;
}
signed main(){
	speed;
	cin >> n;
	for(int i = 0;i<n;i++) cin >> val[i];
	hpar[0] = 0;
	par[0] = -1;
	hld[0].push_back({val[0],1LL});
	for(int i = 0;i<n-1;i++){
		int a,b; cin >> a >> b; a--; b--;
		adj[a].push_back(b);
		edges.push_back({a,b});
	}
	dfs_sz(0);
	dfs_hld(0);
	for(int i = 0;i<n-1;i++){
		chain.clear();
		disc.clear();
		int a = edges[i].fi;
		while(a != -1){
			int apar = hpar[a];
			int siz = depth[a] - depth[apar] + 1;
			vec<pi> temp;
			for(int k = 0;k<sz(hld[apar]);k++){	
				auto [cnt,v] = hld[apar][k];
				if(cnt <= siz){
					siz -= cnt;
					temp.push_back({cnt,v});
				}
				else{
					temp.push_back({siz,v});
					break;
				}
			}
			for(int i = sz(temp)-1;i>=0;i--){
				chain.push_back(temp[i]);
				disc.push_back(temp[i].second);
			}
			a = par[apar];
		}
		reverse(all(chain));
		sort(all(disc));
		disc.erase(unique(all(disc)),disc.end());
		init(sz(disc));
		int ans = 0,mx = 0;	
		for(auto [u,v]: chain){
			v = lower_bound(all(disc),v) - disc.begin();
			ans += u * (mx - qry(v));
			mx += u;
			upd(v,u);
		}
		cout << ans << endl;
		int u = edges[i].se;
		while(u != -1){
			int nxtu = hpar[u];
			int siz = depth[u] - depth[nxtu] + 1;
			while(sz(hld[nxtu]) > 0){
				int cnt = hld[nxtu][0].first;
				if(cnt <= siz){
					siz -= cnt;
					hld[nxtu].pop_front();
				}
				else{
					hld[nxtu][0].first -= siz;
					break;
				}
			}
			hld[nxtu].push_front({depth[u] - depth[nxtu] + 1,val[edges[i].se]});
			u = par[nxtu];
		}
	}
}	
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...