Submission #1241831

#	Time	Username	Problem	Language	Result	Execution time	Memory
1241831		nibert	Tree (IOI24_tree)	C++20	0 / 100	2096 ms	26344 KiB

tree

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

typedef long long ll;

const int MAXN = 200005;

int N;
vector<int> tree[MAXN];
int weight[MAXN];
bool isLeaf[MAXN];

// For Bottom-up
int parent[MAXN];

// For Top-down
struct Segment {
    ll a = 0, b = 0;
};

vector<Segment> segments;

void dfs_mark_leaves(int u) {
    isLeaf[u] = true;
    for (int v : tree[u]) {
        isLeaf[u] = false;
        dfs_mark_leaves(v);
    }
}

ll bottom_up_greedy(int L, int R) {
    vector<int> C(N, 0);
    vector<ll> sub_sum(N);

    for (int i = 0; i < N; ++i) {
        if (isLeaf[i]) C[i] = L;
    }

    for (int i = N - 1; i >= 0; --i) {
        sub_sum[i] += C[i];
        if (parent[i] != -1)
            sub_sum[parent[i]] += sub_sum[i];
    }

    priority_queue<pair<int, int>> pq; // max-heap by weight
    for (int i = 0; i < N; ++i) {
        if (C[i] > 0) pq.emplace(weight[i], i);
    }

    for (int i = N - 1; i >= 0; --i) {
        while (sub_sum[i] > R) {
            while (!pq.empty()) {
                int w = pq.top().first;
                int j = pq.top().second;
                pq.pop();
                
                // check if reducing C[j] by 1 keeps all ancestor subtree sums >= L
                bool valid = true;
                int x = j;
                while (x != -1) {
                    if (sub_sum[x] - 1 < L) {
                        valid = false;
                        break;
                    }
                    x = parent[x];
                }

                if (valid) {
                    C[j]--;
                    x = j;
                    while (x != -1) {
                        sub_sum[x]--;
                        x = parent[x];
                    }
                    if (C[j] > 0) pq.emplace(weight[j], j);
                    break;
                }
            }
        }
    }

    ll cost = 0;
    for (int i = 0; i < N; ++i)
        cost += 1LL * abs(C[i]) * weight[i];
    return cost;
}

int leaf_count[MAXN];
int min_internal_weight;

void dfs_count_leaves(int u) {
    if (isLeaf[u]) {
        leaf_count[u] = 1;
        return;
    }
    leaf_count[u] = 0;
    for (int v : tree[u]) {
        dfs_count_leaves(v);
        leaf_count[u] += leaf_count[v];
    }
}

ll top_down_greedy(int u, int L, int R) {
    if (leaf_count[u] * L <= R) {
        return 1LL * L * leaf_count[u] * weight[u];
    }

    if (isLeaf[u]) return 1LL * L * weight[u];

    ll extra = (1LL * leaf_count[u] * L - R) * weight[u];
    ll res = extra;
    for (int v : tree[u]) {
        res += top_down_greedy(v, L, R);
    }
    return res;
}

void init(vector<int> P, vector<int> W) {
    N = P.size();
    for (int i = 0; i < N; ++i) {
        tree[i].clear();
        isLeaf[i] = false;
    }

    for (int i = 1; i < N; ++i) {
        tree[P[i]].push_back(i);
        parent[i] = P[i];
    }
    parent[0] = -1;

    for (int i = 0; i < N; ++i) weight[i] = W[i];

    dfs_mark_leaves(0);
    dfs_count_leaves(0);

    min_internal_weight = INT_MAX;
    for (int i = 0; i < N; ++i) {
        if (!isLeaf[i]) min_internal_weight = min(min_internal_weight, weight[i]);
    }
}

ll query(int L, int R) {
    if (N <= 2000 || R - L <= 10 || R <= 10 || L == R || N <= 60'000) {
        return bottom_up_greedy(L, R);
    }
    return top_down_greedy(0, L, R);
}

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