#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define debug(x) cerr << #x << " " << x << "\n"
#define debugs(x) cerr << #x << " " << x << " "
#pragma GCC optimize("Ofast,unroll-loops")
using namespace std;
using namespace __gnu_pbds;
typedef long long ll;
typedef pair <ll, ll> pii;
typedef pair <long double, pii> muchie;
typedef tree <ll, null_type, less_equal <ll>, rb_tree_tag, tree_order_statistics_node_update> OST;
const ll NMAX = 100001;
const ll INF = (1LL << 60);
const ll HALF = (1LL << 59);
const ll MOD = 1000000007;
const ll BLOCK = 318;
const ll base = 31;
const ll nr_of_bits = 21;
ll n;
ll dp[NMAX][nr_of_bits + 1];
class All {
public:
ll all[NMAX];
ll lazy[NMAX];
void propaga(ll node, ll st, ll dr) {
if(!lazy[node])
return;
if(st != dr) {
lazy[node * 2] = lazy[node];
lazy[node * 2 + 1] = lazy[node];
}
all[node] = lazy[node];
lazy[node] = 0;
}
void update(ll node, ll st, ll dr, ll a, ll b, ll val) {
propaga(node, st, dr);
if(a <= st && dr <= b) {
lazy[node] = val;
return;
}
ll mid = (st + dr) / 2;
if(a <= mid) {
update(node * 2, st, mid, a, b, val);
}
if(b > mid) {
update(node * 2 + 1, mid + 1, dr, a, b, val);
}
propaga(node * 2, st, mid);
propaga(node * 2 + 1, mid + 1, dr);
all[node] = max(all[node * 2], all[node * 2 + 1]);
}
ll maxim(ll node, ll st, ll dr, ll a, ll b) {
propaga(node, st, dr);
if(a <= st && dr <= b) {
return all[node];
}
ll mid = (st + dr) / 2;
ll sol = 0;
if(a <= mid) {
sol = max(sol, maxim(node * 2, st, mid, a, b));
}
if(b > mid) {
sol = max(sol, maxim(node * 2 + 1, mid + 1, dr, a, b));
}
return sol;
}
} st;
ll poz[NMAX];
ll stamp;
ll sz[NMAX];
ll heavy[NMAX];
ll a[NMAX];
vector <ll> v[NMAX];
void getsize(ll node, ll p) {
sz[node] = 1;
for(auto x : v[node]) {
if(x == p)
continue;
getsize(x, node);
if(heavy[node] == 0 || sz[x] > sz[heavy[node]])
heavy[node] = x;
sz[node] += sz[x];
}
}
ll capat[NMAX];
void DFS(ll node, ll p, ll sf) {
dp[node][0] = p;
capat[node] = sf;
poz[node] = ++stamp;
a[stamp] = node;
if(heavy[node])
DFS(heavy[node], node, sf);
for(auto x : v[node]) {
if(x == p || x == heavy[node])
continue;
DFS(x, node, x);
}
}
pii edges[NMAX];
ll c[NMAX];
void baga(ll node, ll col){
while(node != 0){
st.update(1, 1, n, poz[capat[node]], poz[node], col);
node = dp[capat[node]][0];
}
}
struct ura{
ll first, second, third;
};
ll sol;
ll aib[NMAX];
ll nn;
ura nr[NMAX];
ll cv[NMAX];
void update(ll x, ll val){
for(ll i = x; i <= nn; i+=i&(-i))
aib[i] += val;
}
ll query(ll x){
ll s = 0;
for(ll i = x; i > 0; i-=i&(-i)){
s += aib[i];
}
return s;
}
bool cmp(ura a, ura b){
return a.first < b.first;
}
void calc(){
ll i;
for(i = 0; i <= nn; i++)
aib[i] = 0, cv[i] = 0;
sort(nr + 1, nr + nn + 1, cmp);
ll cnt = 0;
for(i = 1; i <= nn; i++){
if(i == 1 || nr[i].first != nr[i - 1].first){
cnt++;
}
cv[nr[i].second] = cnt;
}
for(i = 1; i <= nn; i++){
sol += (query(cv[i]) * nr[i].third);
update(cv[i], nr[i].third);
}
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
cout.tie(0);
ll i, j;
cin >> n;
for(i = 1; i <= n; i++) {
cin >> c[i];
}
for(i = 1; i < n; i++) {
ll x, y;
cin >> x >> y;
edges[i] = {x, y};
v[x].push_back(y);
v[y].push_back(x);
}
getsize(1, 0);
DFS(1, 0, 1);
for(j = 1; j <= nr_of_bits; j++)
for(i = 1; i <= n; i++)
dp[i][j] = dp[dp[i][j - 1]][j - 1];
st.update(1, 1, n, poz[1], poz[1], c[1]);
for(i = 1; i < n; i++) {
ll x = edges[i].first;
ll y = edges[i].second;
ll node = y;
sol = 0;
nn = 0;
while(node != 1){
ll dist = 1;
node = dp[node][0];
ll col = st.maxim(1, 1, n, poz[node], poz[node]);
ll pas = nr_of_bits;
while(pas >= 0){
ll nxt = dp[node][pas];
if(nxt != 0 && st.maxim(1, 1, n, poz[nxt], poz[nxt]) == col){
node = nxt;
dist += (1 << pas);
}
pas--;
}
nr[++nn] = {col, nn, dist};
}
calc();
baga(y, c[y]);
cout << sol << "\n";
}
return 0;
}
Compilation message
construction.cpp: In function 'int main()':
construction.cpp:191:12: warning: unused variable 'x' [-Wunused-variable]
191 | ll x = edges[i].first;
| ^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
2636 KB |
Output is correct |
| 2 |
Correct |
2 ms |
2636 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
2764 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
2636 KB |
Output is correct |
| 2 |
Correct |
2 ms |
2636 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
2764 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
2636 KB |
Output is correct |
| 2 |
Correct |
2 ms |
2636 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
2764 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
