Submission #1239825

#	Time	Username	Problem	Language	Result	Execution time	Memory
1239825		hugo_pm	Tree (IOI24_tree)	C++20	41 / 100	2097 ms	30536 KiB

tree

#include "tree.h"
#include <cassert>
#include <set>
#define int long long
using i32 = int32_t;
using namespace std;

int N;
vector<int> parent, weight;
vector<vector<int>> children;

void init(vector<i32> P, vector<i32> W) {
    parent = vector<int>(P.begin(), P.end());
    weight = vector<int>(W.begin(), W.end());
    N = (int)parent.size();
    children.resize(N);
    for (int node = 1; node < N; ++node) {
        children[parent[node]].push_back(node);
    }
}

struct Helper {
    int cost, stock;
    auto operator<=>(const Helper&) const = default;
};

class MyPq {
    friend MyPq* merge(MyPq*, MyPq*);
    public:
    void clear() {
        helpers.clear();
        sum_stock = 0;
    }
    void push(int cost, int stock) {
        if (stock == 0) return;
        helpers.insert({cost, stock});
        sum_stock += stock;
    }
    void trim(int lim_sum) {
        if (sum_stock > lim_sum) {
            consume(sum_stock - lim_sum, true);
            assert(sum_stock == lim_sum);
        }
    }
    int cheap_help(int target) {
        return consume(target, false);
    }
    private:
    int consume(int target, bool worst) {
        int cur_cost = 0;
        while (target > 0) {
            assert(!helpers.empty());
            auto it = worst ? prev(helpers.end()) : helpers.begin();
            Helper nxt = *it;
            helpers.erase(it);
            int taking = min(target, nxt.stock);
            target -= taking;
            sum_stock -= taking;
            nxt.stock -= taking;
            cur_cost += taking * nxt.cost; 
            if (nxt.stock)
                helpers.insert(nxt);
        }
        return cur_cost;
    }
    multiset<Helper> helpers;
    int sum_stock = 0;
};

MyPq* merge(MyPq* small, MyPq* large) {
    if (small->helpers.size() > large->helpers.size())
        swap(small, large);
    for (const Helper &h : small->helpers)
        large->push(h.cost, h.stock);
    small->clear();
    return large;
}

int solve(int L, int R) {
    vector<MyPq> block_allocation(N);
    vector<MyPq*> dp(N);
    /// not always up to date!
    vector<int> tmp_sum_subtree(N);
    for (int i = 0; i < N; ++i) {
        dp[i] = &block_allocation[i];
    }
    int answer = 0;
    for (int node = N-1; node >= 0; --node) {
        int &cur_sum = tmp_sum_subtree[node];
        if (children[node].empty()) {
            answer += weight[node] * L;
            cur_sum = L;
            continue;
        }
        MyPq* &cur = dp[node];
        for (int child : children[node]) {
            cur = merge(cur, dp[child]);
            dp[child] = nullptr;
            cur_sum += tmp_sum_subtree[child];
        }
        cur->push(weight[node], cur_sum - L);
        if (cur_sum > R) {
            answer += cur->cheap_help(cur_sum - R);
            cur_sum = R;
        }
        cur->trim(cur_sum - L);
    }
    return answer;
}

long long query(i32 L, i32 R) {
    return solve(L, R);
}

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