Submission #1284784

#	Time	Username	Problem	Language	Result	Execution time	Memory
1284784		icebear	Construction of Highway (JOI18_construction)	C++20	100 / 100	492 ms	42940 KiB

construction

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

typedef long long ll;
typedef pair<int, int> ii;
typedef pair<int, ii> iii;

template<class T>
    bool minimize(T &a, const T &b) {
        if (a > b) return a = b, true;
        return false;
    }

template<class T>
    bool maximize(T &a, const T &b) {
        if (a < b) return a = b, true;
        return false;
    }

#define FOR(i,a,b) for(int i=(a); i<=(b); ++i)
#define FORR(i,a,b) for(int i=(a); i>=(b); --i)
#define REP(i, n) for(int i=0; i<(n); ++i)
#define RED(i, n) for(int i=(n)-1; i>=0; --i)
#define MASK(i) (1LL << (i))
#define BIT(S, i) (((S) >> (i)) & 1)
#define mp make_pair
#define pb push_back
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define task "icebear"
/*END OF TEMPLATE. ICEBEAR AND THE CAT WILL WIN VOI26 */

const int MOD = 1e9 + 7;
const int inf = 1e9 + 27092008;
const ll INF = 1e18 + 27092008;
const int N = 200000 + 5;
struct SegmentTree {
    int n;
    vector<ii> node;
    SegmentTree(int _n = 0): n(_n), node(4 * _n + 5, mp(0, 0)) {}
    void update(int pos, int val) {
        if (pos > n) return;
        int id = 1, l = 1, r = n;
        while(l < r) {
            int mid = (l + r) >> 1;
            if (pos > mid) id = (id << 1 | 1), l = mid + 1;
            else id = (id << 1), r = mid;
        }
        node[id] = mp(val, pos);
        while(id > 1) {
            id >>= 1;
            node[id] = max(node[id << 1], node[id << 1 | 1]);
        }
    }

    ii get(int id, int l, int r, int u, int v) {
        if (l > v || r < u) return mp(0, 0);
        if (u <= l && r <= v) return node[id];
        int mid = (l + r) >> 1;
        return max(get(id << 1, l, mid, u, v), get(id << 1 | 1, mid + 1, r, u, v));
    }

    ii get(int u, int v) {
        return get(1, 1, n, u, v);
    }
} IT[N];
int n, C[N], sz[N], hvy[N], pos[N], par[N], len[N], head[N], curChain = 1, chain[N];
bool active[N];
vector<int> G[N], compress, depos[N];
ii edge[N];
int ft[N];

void update(int x, int v) {
    for(; x <= n; x += x & -x) ft[x] += v;
}

int get(int x) {
    int ans = 0;
    for(; x; x -= x & -x) ans += ft[x];
    return ans;
}

void dfs(int u) {
    sz[u] = 1;
    for(int v : G[u]) if (v != par[u]) {
        par[v] = u;
        dfs(v);
        sz[u] += sz[v];
        if (sz[v] > sz[hvy[u]]) hvy[u] = v;
    }
}

void hld(int u) {
    if (head[curChain] == 0) head[curChain] = u;
    pos[u] = ++len[curChain];
    depos[curChain].pb(u);
    chain[u] = curChain;

    if (hvy[u] > 0) hld(hvy[u]);

    for(int v : G[u]) if (v != par[u] && v != hvy[u]) {
        curChain++;
        hld(v);
    }
}

void init(void) {
    cin >> n;
    FOR(i, 1, n) cin >> C[i], compress.pb(C[i]);
    sort(all(compress));
    compress.resize(unique(all(compress)) - compress.begin());
    FOR(i, 1, n) C[i] = upper_bound(all(compress), C[i]) - compress.begin();
    FOR(i, 2, n) {
        int u, v;
        cin >> u >> v;
        G[u].pb(v);
        G[v].pb(u);
        edge[i] = mp(u, v);
    }
}

void process(void) {
    FOR(i, 1, n) depos[i].pb(0);
    dfs(1);
    hld(1);
    FOR(i, 1, curChain) IT[i] = SegmentTree(len[i]);
    IT[chain[1]].update(pos[1], +1);
    active[1] = true;

    FOR(i, 2, n) {
        vector<ii> updates;
        int u = edge[i].fi, last_chain = -1;
        ll ans = 0;
        while(u > 0) {
            ii tmp = IT[chain[u]].get(1, pos[u]);
            int hd = (tmp.fi == 0 ? head[chain[u]] : depos[chain[u]][tmp.se]);

            ans += (pos[u] - pos[hd] + 1) * get(C[hd] - 1);
            update(C[hd], pos[u] - pos[hd] + 1);
            updates.emplace_back(C[hd], pos[u] - pos[hd] + 1);
            IT[chain[hd]].update(pos[hd], 0);
            if (pos[u] + 1 <= len[chain[u]]) {
                int v = depos[chain[u]][pos[u] + 1];
                if (active[v] && chain[v] != last_chain && IT[chain[hd]].get(pos[v], pos[v]).fi == 0) {
                    IT[chain[hd]].update(pos[v], +1);
                    C[v] = C[hd];
                }
            }

            C[hd] = C[edge[i].se];
            last_chain = chain[u];
            u = par[hd];
        }
        active[edge[i].se] = true;
        for(ii &x : updates) update(x.fi, -x.se);
        cout << ans << '\n';
        cerr << ans << '\n';

//        if (i == 6) {
//            cout << "DEBUG\n";
//            FOR(t, 1, n) cout << t << ' ' << C[t] << ' ' << pos[t] << ' ' << chain[t] << ' ' << len[chain[t]] << ' ' << IT[chain[t]].get(1, 1, n, pos[t], pos[t]).fi << ' ' << IT[chain[t]].get(1, 1, n, pos[t], pos[t]).se << '\n';
//        }
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    if (fopen(task".inp", "r")) {
        freopen(task".inp", "r", stdin);
        freopen(task".out", "w", stdout);
    }
    int tc = 1;
//    cin >> tc;
    while(tc--) {
        init();
        process();
    }
    return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:172:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  172 |         freopen(task".inp", "r", stdin);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
construction.cpp:173:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  173 |         freopen(task".out", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~

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