This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
//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, c[xn], par[xn], hvy[xn], sz[xn], arshad[xn];
int stm[xn], TM, ptr, b[xn], fen[xn], seg[xn << 2];
vector <int> adj[xn], cprs;
vector <pii> vec;
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);
}
void modify(int id, int l, int r, int pos, int val){
if (r - l == 1){
seg[id] = val;
return;
}
int mid = l + r >> 1;
if (pos < mid)
modify(lc, l, mid, pos, val);
else
modify(rc, mid, r, pos, val);
seg[id] = seg[lc];
if (seg[rc] - seg[lc])
seg[id] = - 1;
}
void build(int id, int l, int r){
seg[id] = - 1;
if (r - l == 1)
return;
int mid = l + r >> 1;
build(lc, l, mid), build(rc, mid, r);
}
void solve(int id, int l, int r, int ql, int qr, int val){
if (qr <= l || r <= ql || qr <= ql)
return;
if (ql <= l && r <= qr && seg[id] + 1){
vec.push_back({seg[id], r - l});
seg[id] = val;
return;
}
int mid = l + r >> 1;
if (seg[id] + 1)
seg[lc] = seg[rc] = seg[id];
solve(rc, mid, r, ql, qr, val);
solve(lc, l, mid, ql, qr, val);
seg[id] = seg[lc];
if (seg[rc] - seg[id])
seg[id] = - 1;
if (ql <= l && r <= qr)
seg[id] = val;
}
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);
build(1, 0, TM);
for (int i = 1; i <= n; ++ i)
modify(1, 0, TM, stm[i], c[i]);
for (int i = 1; i < n; ++ i){
int v = b[i], col = c[v];
vec.clear();
v = par[v];
while (v){
solve(1, 0, TM, stm[arshad[v]], stm[v] + 1, col);
v = par[arshad[v]];
}
ll ans = 0;
for (pii x : vec){
ans += 1ll * x.B * gefen(x.A - 1);
mofen(x.A, x.B);
}
cout << ans << endl;
for (pii x : vec)
mofen(x.A, - x.B);
}
return 0;
}
Compilation message (stderr)
construction.cpp: In function 'void modify(int, int, int, int, int)':
construction.cpp:77:14: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
77 | int mid = l + r >> 1;
| ~~^~~
construction.cpp: In function 'void build(int, int, int)':
construction.cpp:90:14: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
90 | int mid = l + r >> 1;
| ~~^~~
construction.cpp: In function 'void solve(int, int, int, int, int, int)':
construction.cpp:101:14: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
101 | int mid = l + r >> 1;
| ~~^~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |