Submission #1258556

#TimeUsernameProblemLanguageResultExecution timeMemory
1258556AHOKAConstruction of Highway (JOI18_construction)C++20
100 / 100
267 ms66368 KiB
#include <bits/stdc++.h>

#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")

using namespace std;
 
#define threesum cin.tie(NULL); cout.tie(NULL); ios_base::sync_with_stdio(false)
#define all(a) a.begin(), a.end()
#define F first
#define S second
#define int long long
#define pii pair<int, int>
#define ppp pair<int, pii>
#define dout cout << fixed << setprecision(15)
#define mid ((l + r) >> 1)
#define lc (id << 1)
#define rc (lc + 1)

const int maxn = 1e6 + 10, maxm = 2e2 + 10, oo = 1e18 + 10, lg = 20, sq = 30, mod = 998244353;

int n, a[maxn], sz[maxn], hvy[maxn], h[maxn], up[maxn], par[maxn];

vector<pii> b;

int fen[maxn];

void add(int i, int val){
    for (; i <= n; i += i & -i)
        fen[i] += val;
}
 
int get(int i){
    int res = 0;
    for (; i > 0; i -= i & -i)
        res += fen[i];
    return res;
}

vector<int> adj[maxn];

vector<pii> stk[maxn];

void pre_dfs(int v, int p){
    h[v] = h[p] + 1;
    par[v] = p;
    sz[v] = 1;
    for(auto u : adj[v])
        if(u ^ p){
            pre_dfs(u, v);
            sz[v] += sz[u];
            if(!hvy[v] || (sz[u] > sz[hvy[v]]))
                hvy[v] = u;
        }
}

void dfs(int v, int p, bool f = 0)
{
    up[v] = (f ? up[p] : v);

    for(auto u : adj[v])
        if(u ^ p)
            dfs(u, v, (u == hvy[v]));
}

pii e[maxn];

void go(int v, int x){
    if(!v)
        return;

    vector<pii> c;
    int p = up[v];
    int last = h[p] - 1;

    while (stk[p].size() && stk[p].back().F <= h[v])
    {
        c.push_back({stk[p].back().S, stk[p].back().F - last});
        last = stk[p].back().F;
        stk[p].pop_back();
    }

    if(stk[p].size())
        c.push_back({stk[p].back().S, h[v] - last});

    stk[p].push_back({h[v], x});

    go(par[p], x);

    for(auto i : c)
        b.push_back(i);
}

int inv(){
    int res = 0;

    vector<int> cmp;
    for(auto [i, cnt] : b)
        cmp.push_back(i);

    sort(all(cmp));
    cmp.resize(unique(all(cmp)) - cmp.begin());

    for (int i = b.size() - 1; i >= 0;i--){
        b[i].F = lower_bound(all(cmp), b[i].F) - cmp.begin() + 1;
        res += b[i].S * get(b[i].F - 1);
        add(b[i].F, b[i].S);
    }

    for(auto [i, cnt] : b)
        add(i, -cnt);

    return res;
}

void solve()
{
    cin >> n;
    for (int i = 1; i <= n;i++)
        cin >> a[i];

    for (int i = 1; i < n;i++){
        int u, v;
        cin >> u >> v;
        e[i] = {u, v};
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    pre_dfs(1, 0);
    dfs(1, 0);

    stk[1].push_back({h[1], a[1]});
    for (int i = 1; i < n; i++)
    {
        int v = e[i].S;

        b.clear();

        go(par[v], a[v]);

        stk[up[v]] = {{h[v], a[v]}};

        cout << inv() << "\n";
    }
}

signed main()
{
    threesum;
    int t = 1;
    //cin >> t;
    while(t--)
        solve();
}
/*
3
1 2 3
1 2
2 3
*/
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...