Submission #516258

#	Time	Username	Problem	Language	Result	Execution time	Memory
516258		AmShZ	Construction of Highway (JOI18_construction)	C++11	100 / 100	372 ms	90704 KiB

construction

//khodaya khodet komak kon
# include <bits/stdc++.h>
 
using namespace std;
 
typedef long long                                        ll;
typedef long double                                      ld;
typedef pair <int, int>                                  pii;
typedef pair <pii, int>                                  ppi;
typedef pair <int, pii>                                  pip;
typedef pair <pii, pii>                                  ppp;
typedef pair <ll, ll>                                    pll;
 
# define A                                               first
# define B                                               second
# define endl                                            '\n'
# define sep                                             ' '
# define all(x)                                          x.begin(), x.end()
# define kill(x)                                         return cout << x << endl, 0
# define SZ(x)                                           int(x.size())
# define lc                                              id << 1
# define rc                                              id << 1 | 1
# define fast_io                                         ios::sync_with_stdio(0);cin.tie(0); cout.tie(0);
 
ll power(ll a, ll b, ll md) {return (!b ? 1 : (b & 1 ? a * power(a * a % md, b / 2, md) % md : power(a * a % md, b / 2, md) % md));}
 
const int xn = 1e5 + 10;
const int xm = - 20 + 10;
const int sq = 320;
const int inf = 1e9 + 10;
const ll INF = 1e18 + 10;
const ld eps = 1e-15;
const int mod = 998244353;
const int base = 257;
 
int n, m, c[xn], par[xn], hvy[xn], sz[xn], arshad[xn];
int stm[xn], TM, ptr, b[xn], fen[xn];
pii a[xn];
stack <ppi> sk[xn];
vector <int> adj[xn], cprs;
map <int, int> mp;
 
void mofen(int pos, int val){
	for (pos += 5; pos < xn; pos += pos & - pos)
		fen[pos] += val;
}
int gefen(int pos){
	int res = 0;
	for (pos += 5; 0 < pos; pos -= pos & - pos)
		res += fen[pos];
	return res;
}
void preDFS(int v){
	sz[v] = 1;
	for (int u : adj[v]){
		preDFS(u);
		sz[v] += sz[u];
		if (sz[hvy[v]] < sz[u])
			hvy[v] = u;
	}
}
void DFS(int v){
	stm[v] = TM ++;
	arshad[v] = v;
	if (v == hvy[par[v]])
		arshad[v] = arshad[par[v]];
	if (hvy[v])
		DFS(hvy[v]);
	for (int u : adj[v])
		if (u != hvy[v])
			DFS(u);
	sk[arshad[v]].push({{stm[v], stm[v]}, c[v]});
}
void add(pii x){
	if (m && a[m - 1].A == x.A)
		a[m - 1].B += x.B;
	else
		a[m ++] = x;
}
void solve(int id, int x, int col){
	vector <pii> vec;
	while (SZ(sk[id]) && sk[id].top().A.B <= x){
		vec.push_back({sk[id].top().B, sk[id].top().A.B - sk[id].top().A.A + 1});
		sk[id].pop();
	}
	if (SZ(sk[id]) && sk[id].top().A.A <= x){
		vec.push_back({sk[id].top().B, x - sk[id].top().A.A + 1});
		sk[id].top().A.A = x + 1;
	}
	if (SZ(sk[id]) && sk[id].top().B == col)
		sk[id].top().A.A = stm[id];
	else
		sk[id].push({{stm[id], x}, col});
	reverse(all(vec));
	for (pii x : vec)
		add(x);
}
 
int main(){
	fast_io;
 
	cin >> n;
	for (int i = 1; i <= n; ++ i)
		cin >> c[i], cprs.push_back(c[i]);
	sort(all(cprs));
	for (int i = 0; i < SZ(cprs); ++ i)
		if (!i || cprs[i] != cprs[i - 1])
			mp[cprs[i]] = ptr ++;
	for (int i = 1; i <= n; ++ i)
		c[i] = mp[c[i]];
	for (int i = 1; i < n; ++ i){
		int v, u;
		cin >> v >> u;
		par[u] = v;
		adj[v].push_back(u);
		b[i] = u;
	}
	preDFS(1);
	DFS(1);
	for (int i = 1; i < n; ++ i){
		int v = b[i], col = c[v];
		v = par[v], m = 0;
		while (v){
			solve(arshad[v], stm[v], col);
			v = par[arshad[v]];
		}
		ll ans = 0;
		for (int i = 0; i < m; ++ i){
			ans += 1ll * a[i].B * gefen(a[i].A - 1);
			mofen(a[i].A, a[i].B);
		}
		cout << ans << endl;
		for (int i = 0; i < m; ++ i)
			mofen(a[i].A, - a[i].B);
	}
 
	return 0;
}

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