This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
// Knapsack DP is harder than FFT.
#include<bits/stdc++.h>
using namespace std;
typedef long long ll; typedef pair<int,int> pii;
#define ff first
#define ss second
#define pb emplace_back
#define AI(x) begin(x),end(x)
#ifdef OWO
#define debug(args...) SDF(#args, args)
#define OIU(args...) ostream& operator<<(ostream&O,args)
#define LKJ(S,B,E,F) template<class...T>OIU(S<T...>s){O<<B;int c=0;for(auto i:s)O<<(c++?", ":"")<<F;return O<<E;}
LKJ(vector,'[',']',i)LKJ(deque,'[',']',i)LKJ(set,'{','}',i)LKJ(multiset,'{','}',i)LKJ(unordered_set,'{','}',i)LKJ(map,'{','}',i.ff<<':'<<i.ss)LKJ(unordered_map,'{','}',i.ff<<':'<<i.ss)
template<class...T>void SDF(const char* s,T...a){int c=sizeof...(T);if(!c){cerr<<"\033[1;32mvoid\033[0m\n";return;}(cerr<<"\033[1;32m("<<s<<") = (",...,(cerr<<a<<(--c?", ":")\033[0m\n")));}
template<class T,size_t N>OIU(array<T,N>a){return O<<vector<T>(AI(a));}template<class...T>OIU(pair<T...>p){return O<<'('<<p.ff<<','<<p.ss<<')';}template<class...T>OIU(tuple<T...>t){return O<<'(',apply([&O](T...s){int c=0;(...,(O<<(c++?", ":"")<<s));},t),O<<')';}
#else
#define debug(...) ((void)0)
#endif
struct LISAN : vector<int> {
void done() { sort(AI()); erase(unique(AI()), end()); }
int operator () (const int& v) const { return lower_bound(AI(), v) - begin(); }
};
struct BIT {
int n; vector<int> val;
BIT(int _n = 0): n(_n), val(n + 1, 0) {}
void modify(int p, int v) {
for (; p <= n; p += p & -p)
val[p] += v;
}
int query(int p) {
int r = 0;
for (; p > 0; p -= p & -p)
r += val[p];
return r;
}
};
const int kN = 100'001;
int N, C[kN], A[kN], B[kN];
int par[kN];
vector<int> chi[kN];
int dep[kN], siz[kN];
void dfs_siz(int u) {
siz[u] = 1;
for (int& v: chi[u]) {
dep[v] = dep[u] + 1;
dfs_siz(v);
siz[u] += siz[v];
if (siz[v] > siz[chi[u][0]])
swap(v, chi[u][0]);
}
}
vector<pii> chain[kN];
int head[kN];
void dfs_head(int u, int h) {
head[u] = h;
for (const int& v: chi[u])
dfs_head(v, v == chi[u][0] ? h : v);
if (chain[h].empty())
chain[h].pb(-1, dep[u] + 1);
if (chain[h].back().ff != C[u])
chain[h].pb(C[u], dep[u]);
else
chain[h].back().ss = dep[u];
}
signed main() {
ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
cin >> N;
LISAN lisan;
for (int i = 1; i <= N; ++i) {
cin >> C[i];
lisan.pb(C[i]);
}
lisan.done();
for (int i = 1; i <= N; ++i)
C[i] = lisan(C[i]) + 1;
for (int i = 1; i < N; ++i) {
cin >> A[i] >> B[i];
par[B[i]] = A[i];
chi[A[i]].pb(B[i]);
}
dfs_siz(1);
dfs_head(1, 1);
BIT bit(lisan.size());
vector<pii> vals, tmp;
for (int i = 1; i < N; ++i) {
int u = A[i];
do {
int h = head[u];
while (8e7 /* is stronger than me */ ) {
pii cur = chain[h].back();
if (cur.ss > dep[u]) break;
chain[h].pop_back();
if (chain[h].back().ss > dep[u]) {
tmp.pb(cur.ff, dep[u] - cur.ss + 1);
if (chain[h].back().ss > dep[u] + 1)
chain[h].pb(cur.ff, dep[u] + 1);
break;
}
tmp.pb(cur.ff, chain[h].back().ss - cur.ss);
}
while (!tmp.empty()) {
vals.push_back(tmp.back());
tmp.pop_back();
}
u = par[h];
} while (u != 0);
ll ans = 0;
for (const auto& [val, num]: vals) {
ans += 1ll * bit.query(val - 1) * num;
bit.modify(val, num);
}
cout << ans << '\n';
for (const auto& [val, num]: vals)
bit.modify(val, -num);
vals.clear();
u = A[i];
do {
int h = head[u];
chain[h].pb(C[B[i]], dep[h]);
u = par[h];
} while (u != 0);
}
return 0;
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |