Submission #1310535

#TimeUsernameProblemLanguageResultExecution timeMemory
1310535vlomaczkConstruction of Highway (JOI18_construction)C++20
0 / 100
10 ms14720 KiB
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
typedef long long ll;
using namespace __gnu_pbds;
using namespace std;

template <typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

ll M = 100'010;
vector<vector<ll>> g(M);
vector<ll> s(M), d(M), par(M), czy_heavy(M), pre(M), post(M), path_end(M);
vector<vector<pair<ll, ll>>> paths(M);
ll n;

void sajz_dfs(ll v, ll p) {
	s[v] = 1;
	d[v] = d[p]+1;
	par[v] = p;
	for(auto u : g[v]) {
		if(u==p) continue;
		sajz_dfs(u,v);
		s[v] += s[u];
	}
}

void make_hld(ll v, ll p) {
	pair<ll, ll> heavy = {-1, -1};
	for(auto u : g[v]) if(u!=p) heavy = max(heavy, {s[u], u});
	if(heavy.first!=-1) { 
		czy_heavy[v] = 1;
		path_end[heavy.second] = path_end[v];
		make_hld(heavy.second, v);
	}
	for(auto u : g[v]) {
		if(u!=p && u!=heavy.second) {
			path_end[u] = u;
			make_hld(u,v);
		}
	}
}

struct SegTree {
	ll base = 1<<18;
	vector<ll> T;

	SegTree() {
		T.assign(2*base,0);
	}

	void update(ll x, ll val) {
		x += base;
		T[x] += val;
		x/=2;
		while(x>0) {
			T[x] = T[x+x] + T[x+x+1];
			x/=2;
		}
	}

	ll query(ll a, ll b) {
		a += base - 1;
		b += base + 1;
		ll res = 0;
		while(a/2 != b/2) {
			if(a%2==0) res += T[a+1];
			if(b%2==1) res += T[b-1];
			a/=2; b/=2;
		}
		return res;
	}
};

ll query(ll x) {
	vector<pair<ll, ll>> pary;
	while(x) {
		vector<pair<ll, ll>> to_add;
		ll pref = d[path_end[x]];
		for(auto[a,b] : paths[path_end[x]]) {
			if(pref + b > d[x]) {
				to_add.push_back({a, d[x] - pref + 1});
				break;
			} else to_add.push_back({a,b});
			pref += b;
		}
		reverse(paths[path_end[x]].begin(), paths[path_end[x]].end());
		for(auto p : to_add) {
			if(paths[path_end[x]].back()==p) paths[path_end[x]].pop_back();
			else paths[path_end[x]].back().second -= p.second;
		}
		reverse(paths[path_end[x]].begin(), paths[path_end[x]].end());
		reverse(to_add.begin(), to_add.end());
		for(auto p : to_add) pary.push_back(p);
		x = par[path_end[x]];
	}
	reverse(pary.begin(), pary.end());
	ll res = 0;
	SegTree st;
	for(auto[a,b] : pary) {
		res += b * st.query(0,n+10);
		st.update(a,b);
	}
	return res;
}

void update(ll x, ll val) {
	while(x) {
		reverse(paths[path_end[x]].begin(), paths[path_end[x]].end());
		paths[path_end[x]].push_back({val, d[x]-d[path_end[x]] + 1});
		reverse(paths[path_end[x]].begin(), paths[path_end[x]].end());
		x = par[path_end[x]];
	}
}

void scale(vector<ll> &x) {
	map<ll, ll> mapa;
	for(auto i : x) mapa[i]++;
	int nxt = 1;
	for(auto&[a,b] : mapa) b = nxt++; 
	for(auto &i : x) i = mapa[i];
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0); cout.tie(0);

	cin >> n;
	vector<ll> C(n+1);
	for(ll i=1; i<=n; ++i) cin >> C[i];
	scale(C);
	vector<pair<ll, ll>> edges;
	for(ll i=1; i<n; ++i) {
		ll a, b;
		cin >> a >> b;
		g[a].push_back(b);
		g[b].push_back(a);
		edges.push_back({a,b});
	}
	sajz_dfs(1,0);
	path_end[1] = 1;
	make_hld(1,0);
	for(auto[a,b] : edges) {
		cout << query(a) << "\n";
		update(b, C[b]);
	}

	return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...