Submission #1272306

#	Time	Username	Problem	Language	Result	Execution time	Memory
1272306		velislavgarkov	Tree (IOI24_tree)	C++20	59 / 100	2094 ms	92180 KiB

tree

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

#define ll long long
#define pli pair <long long, int>
const int MAXN = 2e5 + 10;
const ll INF = 2e9;

struct SegTree {
    int n;
    vector <ll> tree, lazy;
    void init(int N) {
        n = N;
        tree = vector <ll> (4*n, 0); lazy = vector <ll> (4*n, 0);
    }

    void push_lazy(int node, int l, int r) {
        if (lazy[node] == 0) return;
        tree[node] += lazy[node];
        if (l != r) {
            lazy[node * 2] += lazy[node];
            lazy[node * 2 + 1] += lazy[node];
        }
        lazy[node] = 0;
    }

    void upd(int node, int l, int r, int ql, int qr, ll ch) {
        push_lazy(node, l, r);
        if (ql > r || qr < l) return;
        if (l >= ql && r <= qr) {
            lazy[node] += ch;
            push_lazy(node, l, r);
            return;
        }

        int mid = (l + r) / 2;
        upd(node * 2, l, mid, ql, qr, ch);
        upd(node * 2 + 1, mid + 1, r, ql, qr, ch);
        tree[node] = min(tree[node * 2], tree[node * 2 + 1]);
    }

    ll qry(int node, int l, int r, int ql, int qr) {
        push_lazy(node, l, r);
        if (ql > r || qr < l) return INF;
        if (l >= ql && r <= qr) return tree[node];
        int mid = (l + r) / 2;
        return min(qry(node * 2, l, mid, ql, qr), qry(node * 2 + 1, mid + 1, r, ql, qr));
    }

    ll query(int ql, int qr) {
        if (ql > qr) return INF;
        return qry(1, 0, n - 1, ql, qr);
    }
    void update(int ql, int qr, ll ch) {
        if (ql > qr) return;
        upd(1, 0, n - 1, ql, qr, ch);
    }

} tree;

int n;
vector<int> p, w, v[MAXN];
ll sum[MAXN];

int sz[MAXN], nxt[MAXN], head[MAXN], pos[MAXN], cur_pos;

bool ones;
ll cnt[MAXN], leafs[MAXN], suff[MAXN];
ll cnt_leafs, cnt_trees;

void dfs_sz(int x) {
    sz[x] = 1; nxt[x] = -1;
    for (auto i : v[x]) {
        dfs_sz(i);
        sz[x] += sz[i];
        if (nxt[x] == -1 || sz[nxt[x]] < sz[i]) nxt[x] = i;
    }
}

void decompose(int x, int h) {
    head[x] = h;
    pos[x] = cur_pos++;
    if (nxt[x] != -1) decompose(nxt[x], h);

    for (auto i : v[x]) {
        if (i == nxt[x]) continue;
        decompose(i, i);
    }
}

void dfs_ones(int x) {
    if (v[x].empty()) {
        cnt[x] = 1;
        if (w[x] == 1) cnt_leafs++;
        return;
    }

    for (auto i : v[x]) {
        dfs_ones(i);
        cnt[x] += cnt[i];
    }
    if (w[x] == 1 && (p[x] == -1 || w[p[x]] == 0)) {
        leafs[cnt_trees++] = cnt[x];
    }
    if (w[x] == 0) cnt[x] = 1;
}

void init(std::vector<int> P, std::vector<int> W) {
    p = P;
    w = W;
    n = (int)p.size();

    for (int i = 1; i < n; i++) {
        v[p[i]].push_back(i);
    }
    ones = true;
    for (int i = 0; i < n; i++) {
        if (w[i] != 0 && w[i] != 1) ones = false;
    }

    if (ones) {
        cnt_leafs = 0;
        dfs_ones(0);
        sort(leafs, leafs + cnt_trees);
        suff[cnt_trees] = 0;
        for (int i = cnt_trees - 1; i >= 0; i--) {
            suff[i] = suff[i + 1] + leafs[i];
        } 
        return;
    }

    dfs_sz(0);
    cur_pos = 0;
    decompose(0, 0);
}

ll getmin(int u, int ch) {
    ll res = INF;
    while (head[ch] != head[u]) {
        res = min(res, tree.query(pos[head[ch]], pos[ch]));
        ch = p[head[ch]];
    }
    res = min(res, tree.query(pos[u] + 1, pos[ch]));
    return res;
}

void update_path(int u, int ch, ll k) {
    while (head[ch] != head[u]) {
        tree.update(pos[head[ch]], pos[ch], k);
        ch = p[head[ch]];
    }
    tree.update(pos[u] + 1, pos[ch], k);
}

set <pli> s[MAXN];

void merge(int x, int y) {
    if (s[x].size() < s[y].size()) swap(s[x], s[y]);
    for (auto i : s[y]) s[x].insert(i);
}

int find_idx(ll L, ll R) {
    int l, r, mid;
    l = 0; r = cnt_trees - 1;

    while (l <= r) {
        mid = (l + r) / 2;
        if (leafs[mid] * L <= R) l = mid + 1;
        else r = mid - 1;
    }

    return l;
}

long long query(int l, int r) {
    if (ones) {
        ll ans = 0;
        ans += cnt_leafs * (ll)l;
        ll idx = find_idx(l, r);
        ans += max((ll)0, suff[idx] * l - (ll)(cnt_trees - idx) * r);
        return ans;
    }

    ll ans = 0;
    for (int i = 0; i < n; i++) s[i].clear();
    tree.init(n);

    for (int u = n - 1; u >= 0; u--) {
        if (sz[u] == 1) {
            ans += l * (ll)w[u];
            sum[u] = l;
            continue;
        }

        s[u].insert({w[u], u});
        sum[u] = 0;
        for (auto i : v[u]) {
            merge(u, i);
            sum[u] += sum[i];
        }

        while (sum[u] > r) {
            int ch = s[u].begin() -> second;
            if (ch == u) {
                ans += (sum[u] - r) * (ll)w[u];
                sum[u] = r;
                break;
            }

            ll t = getmin(u, ch);
            ll k = min(sum[u] - r, t);
            if (k == 0) {
                s[u].erase(s[u].begin());
                continue;
            }

            update_path(u, ch, -k);
            sum[u] -= k;
            ans += k * (ll)w[ch];

            if (t == k) {
                s[u].erase(s[u].begin());
            }
        }
        tree.update(pos[u], pos[u], sum[u] - l);
    }

    return ans;
}

• Subtask 0: 0 / 0
• Subtask 1: 10 / 10
• Subtask 2: 13 / 13
• Subtask 3: 18 / 18
• Subtask 4: 7 / 7
• Subtask 5: 11 / 11
• Subtask 6: 0 / 22
• Subtask 7: 0 / 19