Submission #915299

#	Time	Username	Problem	Language	Result	Execution time	Memory
915299		Amirreza_Fakhri	Construction of Highway (JOI18_construction)	C++17	100 / 100	221 ms	52048 KiB

construction

// In the name of the God
#include <bits/stdc++.h>
#define ll long long
#define int long long
#define pb push_back
#define F first
#define S second
#define mp make_pair
#define pii pair <int, int>
#define smin(x, y) (x) = min((x), (y))
#define smax(x, y) (x) = max((x), (y))
#define all(x) (x).begin(), (x).end()
using namespace std;

const int inf = 1e9+7;
const int mod = 998244353;
const int maxn = 5e5+5;

int n, c[maxn], par[maxn], he[maxn], node[maxn], heavy[maxn], head[maxn], fen[maxn];
vector <int> uni, adj[maxn];
vector <pii> ch[maxn], vec1, vec2;

int dfs1(int v = 0) {
    int sz = 1, mx = 0;
    heavy[v] = -1;
    for (int u : adj[v]) {
        int csz = dfs1(u);
        sz += csz;
        if (csz > mx) mx = csz, heavy[v] = u;
    }
    return sz;
}

void dfs2(int v = 0) {
    for (int u : adj[v]) {
        if (u == heavy[v]) head[u] = head[v];
        else head[u] = u;
        dfs2(u);
    }
}

int get(int i) {
    int ans = 0;
    for (i++; i; i -= i&-i) ans += fen[i-1];
    return ans;
}

void upd(int i, int x) {
    for (i++; i <= n; i += i&-i) fen[i-1] += x;
}

int processs(int v) {
    int ans = 0;
    while (v != -1) {
        int last = he[head[v]]-1;
        for (int i = ch[head[v]].size()-1; i >= 0; i--) {
            vec1.pb(mp(min(ch[head[v]][i].F, he[v])-last, ch[head[v]][i].S));
            last = ch[head[v]][i].F;
            if (last >= he[v]) break;
        }
        for (int i = vec1.size()-1; i >= 0; i--) {
            // cout << vec1[i].S << ' ' << vec1[i].F << " FF\n";
            vec2.pb(vec1[i]);
            ans += get(vec1[i].S-1)*vec1[i].F;
            upd(vec1[i].S, vec1[i].F);
        }
        vec1.clear();
        v = par[head[v]];
    }
    for (int i = 0; i < vec2.size(); i++) upd(vec2[i].S, -vec2[i].F);
    vec2.clear();
    return ans;
}

void amir(int v) {
    int x = c[v];
    while (v != -1) {
        // cout << v << "\n";
        for (int i = ch[head[v]].size()-1; i >= 0; i--) {
            if (ch[head[v]][i].F > he[v]) break;
            ch[head[v]].pop_back();
        }
        ch[head[v]].pb(mp(he[v], x));
        v = par[head[v]];
    }
    // for (int i = 0; i < ch[0].size(); i++) {
    //     cout << ch[0].size() << ' ' << ch[0][i].F << ' ' << ch[0][i].S << '\n';
    // }
    // cout << "_____\n";
}

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n;
    for (int i = 0; i < n; i++) {
        cin >> c[i];
        uni.pb(c[i]);
    }
    sort(all(uni)); uni.resize(unique(all(uni))-uni.begin());
    for (int i = 0; i < n; i++) c[i] = lower_bound(all(uni), c[i])-uni.begin();
    par[0] = -1;
    for (int i = 0; i < n-1; i++) {
        int v, u; cin >> v >> u;
        adj[--v].pb(--u);
        par[u] = v;
        he[u] = he[v]+1;
        node[i] = u;
    }
    dfs1(), dfs2();
    ch[0].pb(mp(0, c[0]));
    for (int i = 0; i < n-1; i++) {
        cout << processs(par[node[i]]) << '\n';
        amir(node[i]);
        // for (int i = 0
    }
    // cout << processs(2) << '\n';
    // cout << heavy[1] << '\n';
    return 0;
}


/* 
5
1 2 3 4 5
1 2
2 3
2 4
3 5 
*/

Compilation message (stderr)

construction.cpp: In function 'long long int processs(long long int)':
construction.cpp:70:23: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   70 |     for (int i = 0; i < vec2.size(); i++) upd(vec2[i].S, -vec2[i].F);
      |                     ~~^~~~~~~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84