Submission #985592

#	Time	Username	Problem	Language	Result	Execution time	Memory
985592		alextodoran	Construction of Highway (JOI18_construction)	C++17	100 / 100	216 ms	29700 KiB

construction

/**
 _  _   __  _ _ _  _  _ _
 |a  ||t  ||o    d | |o  |
| __    _| | _ | __|  _ |
| __ |/_  | __  /__\ / _\|

**/

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int N_MAX = 100000;

int N;
int C[N_MAX + 2];
int A[N_MAX + 2], B[N_MAX + 2];
vector <int> adj[N_MAX + 2];

int order[N_MAX + 2];
void compress () {
    iota(order + 1, order + N + 1, 1);
    sort(order + 1, order + N + 1, [&] (const int &u, const int &v) {
        return C[u] < C[v];
    });
    int curr = 0, last = 0;
    for (int i = 1; i <= N; i++) {
        if (C[order[i]] > last) {
            last = C[order[i]];
            curr++;
        }
        C[order[i]] = curr;
    }
}

int parent[N_MAX + 2];
int sub[N_MAX + 2];

void dfs (int u) {
    sub[u] = 1;
    for (int v : adj[u]) {
        if (v != parent[u]) {
            parent[v] = u;
            dfs(v);
            sub[u] += sub[v];
        }
    }
}

int cnt_chains;
vector <int> chain[N_MAX + 2];
int chain_id[N_MAX + 2];
int chain_pos[N_MAX + 2];

void hld (int u) {
    chain_pos[u] = (int) chain[cnt_chains].size();
    chain[cnt_chains].push_back(u);
    chain_id[u] = cnt_chains;
    int heavy = 0;
    for (int v : adj[u]) {
        if (v != parent[u] && sub[v] > sub[heavy]) {
            heavy = v;
        }
    }
    if (heavy != 0) {
        hld(heavy);
        for (int v : adj[u]) {
            if (v != parent[u] && v != heavy) {
                cnt_chains++;
                hld(v);
            }
        }
    }
}

struct SegInfo {
    int mn, mx;
    int lazy;
    void update (int val) {
        mn = mx = val;
        lazy = val;
    }
};
SegInfo operator + (const SegInfo &A, const SegInfo &B) {
    return SegInfo{min(A.mn, B.mn), max(A.mx, B.mx), 0};
}

vector <SegInfo> seg_tree[N_MAX + 2];

void build () {
    for (int c = 1; c <= cnt_chains; c++) {
        seg_tree[c].resize((int) chain[c].size() * 4 + 2);
    }
}

void update (int c, int node, int l, int r, int i, int val) {
    if (r <= i) {
        seg_tree[c][node].update(val);
        return;
    }
    int mid = (l + r) / 2;
    int left = node * 2, right = node * 2 + 1;
    if (seg_tree[c][node].lazy != 0) {
        seg_tree[c][left].update(seg_tree[c][node].lazy);
        seg_tree[c][right].update(seg_tree[c][node].lazy);
        seg_tree[c][node].lazy = 0;
    }
    update(c, left, l, mid, i, val);
    if (mid + 1 <= i) {
        update(c, right, mid + 1, r, i, val);
    }
    seg_tree[c][node] = seg_tree[c][left] + seg_tree[c][right];
}
void update (int c, int i, int val) {
    update(c, 1, 0, (int) chain[c].size() - 1, i, val);
}

vector <pair <int, int>> blocks;
void add_block (pair <int, int> block) {
    if (blocks.empty() == false && blocks.back().first == block.first) {
        blocks.back().second += block.second;
    } else {
        blocks.push_back(block);
    }
}
void query (int c, int node, int l, int r, int i) {
    if (r <= i && seg_tree[c][node].mn == seg_tree[c][node].mx) {
        add_block(make_pair(seg_tree[c][node].mn, r - l + 1));
        return;
    }
    int mid = (l + r) / 2;
    int left = node * 2, right = node * 2 + 1;
    if (seg_tree[c][node].lazy != 0) {
        seg_tree[c][left].update(seg_tree[c][node].lazy);
        seg_tree[c][right].update(seg_tree[c][node].lazy);
        seg_tree[c][node].lazy = 0;
    }
    if (mid + 1 <= i) {
        query(c, right, mid + 1, r, i);
    }
    query(c, left, l, mid, i);
}
void query (int c, int i) {
    query(c, 1, 0, (int) chain[c].size() - 1, i);
}

int Fen[N_MAX + 2];
void update_Fen (int pos, int val) {
    for (int i = pos; i <= N; i += i & -i) {
        Fen[i] += val;
    }
}
int query_Fen (int pos) {
    int sum = 0;
    for (int i = pos; i >= 1; i -= i & -i) {
        sum += Fen[i];
    }
    return sum;
}

ll solve (int start) {
//    cout << "solve(" << start << "): ";
    int u = start;
    while (u != 0) {
        int c = chain_id[u];
//        cout << "Chain #" << c << " (" << chain[c][0] << ".." << u << ") [" << chain_pos[u] << "] ";
//        cout << "C = " << C[start] << "\n";
        if (chain_pos[u] - (u == start) >= 0) {
            query(c, chain_pos[u] - (u == start));
        }
        update(c, chain_pos[u], C[start]);
        u = parent[chain[c][0]];
    }
    ll answer = 0;
    for (pair <int, int> block : blocks) {
//        cout << "Block " << block.first << " of size " << block.second << "\n";
        answer += (ll) block.second * query_Fen(block.first - 1);
        update_Fen(block.first, block.second);
    }
    for (pair <int, int> block : blocks) {
        update_Fen(block.first, -block.second);
    }
    blocks.clear();
    return answer;
}

int main () {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    cin >> N;
    for (int u = 1; u <= N; u++) {
        cin >> C[u];
    }
    for (int i = 1; i <= N - 1; i++) {
        cin >> A[i] >> B[i];
        adj[A[i]].push_back(B[i]);
        adj[B[i]].push_back(A[i]);
    }
    compress();
    dfs(1);
    cnt_chains = 1; hld(1);
    build();
    solve(1);
    for (int i = 1; i <= N - 1; i++) {
        cout << solve(B[i]) << "\n";
    }

    return 0;
}

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