Indian IOI Training Camp: Zorgian Amusement Park

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

#define pb push_back
#define mp make_pair
#define f first

#define FOR(i, a, b) for (int i=a; i<b; i++)
#define F0R(i, a) FOR(i, 0, a)

const int MAX = 200005;
const int MOD = 1000000007;
const int MAGIC = 3;

vi g[MAX];
int h[MAX], s[MAX], sz = 0;
// s is the stack
ll tot = 0;
void dfs(int u, int p) {
    // cout << "-----------------------" << u << "---------------------------" << endl;
    int lo = 0, hi = sz-1, szorig = sz, eleorig;
    if (u == 1) { s[sz++] = h[u]; }
    else {
        while (lo < hi) {
            int mid = (lo + hi) / 2;
            if (s[mid] <= h[u]) { hi = mid; }
            else { lo = mid + 1; }
        }
        if (h[u] < s[sz-1]) { lo = sz; }
        if (lo == 0) { tot += sz; }
        else if (h[u] == s[lo]) { tot += sz - lo; }
        else { tot += sz - lo + 1; }
        eleorig = s[lo];
        sz = lo + 1; s[lo] = h[u];
    }
    // cout << "current stack: " << endl;
    // F0R(i, sz) { cout << s[i] << " "; }
    // cout << endl;
    for (int v : g[u]) {
        if (v == p) { continue; }
        dfs(v, u);
    }
    if (u == 1) { return; }
    sz = szorig; s[lo] = eleorig;
}
int main() {
    int n; scanf("%d", &n);
    F0R(i, n-1) {
        int u, v; scanf("%d %d", &u, &v);
        g[u].pb(v); g[v].pb(u);
    }
    vector<int> comp;
    FOR(i, 1, n+1) { scanf("%d", &h[i]); comp.pb(h[i]); }
    sort(comp.begin(), comp.end());
    auto it = unique(comp.begin(), comp.end()); comp.erase(it, comp.end());
    FOR(i, 1, n+1) { h[i] = lower_bound(comp.begin(), comp.end(), h[i]) - comp.begin() + 1; }
    // FOR(i, 1, n+1) { cout << h[i] << endl; }
    dfs(1, 0);
    printf("%lld\n", tot);
}