Submission #1241110

#	Time	Username	Problem	Language	Result	Execution time	Memory
1241110		hugo_pm	Tree (IOI24_tree)	C++20	59 / 100	2096 ms	30772 KiB

tree

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

struct MyVec {
    vector<int> vals, suff_sum; int N;
    MyVec(const vector<int> &_v) : vals(_v), N(_v.size()) {
        sort(vals.begin(), vals.end());
        suff_sum.resize(vals.size() + 1);
        for (int i = N-1; i >= 0; --i) {
            suff_sum[i] = vals[i] + suff_sum[i+1];
        }
    }
    /// returns {sum, count} of values >= target
    pair<int, int> sum_geq(int target) {
        int i = lower_bound(vals.begin(), vals.end(), target) - vals.begin();
        return {suff_sum[i], N - i};
    }
};

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

namespace St45 {
    bool active = true; // set to false in init
    int sum_weight_real_leaves = 0;
    MyVec tool({});
    void prepare() {
        vector<bool> seen(N);
        vector<int> sub_leaves;
        for (int subroot = 0; subroot < N; ++subroot) {
            if (children[subroot].empty())
                sum_weight_real_leaves += weight[subroot];
            if (weight[subroot] > 0 && !seen[subroot]) {
                sub_leaves.push_back(0);
                vector<int> stk { subroot };
                while (!stk.empty()) {
                    int cur = stk.back();
                    stk.pop_back();
                    seen[cur] = true;
                    if (weight[cur] == 0 || children[cur].empty()) {
                        ++sub_leaves.back();
                    } else {
                        for (int child : children[cur]) {
                            stk.push_back(child);
                        }
                    }
                }
            }
        }
        tool = MyVec(sub_leaves);
    }
    int solve(int L, int R) {
        // int ans = L*real_leaves;
        // for (int x : tool.vals)
        //     ans += max(L*x - R, 0ll);
        // return ans;
        auto [sum_above, cnt_above] = tool.sum_geq((R+L-1)/L); // ?
        return L*sum_weight_real_leaves + L*sum_above - R*cnt_above;
    }
}

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);
    }
    for (int w : weight) {
        if (w > 1) St45::active = false;
    }
    if (St45::active) St45::prepare();
}

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) {
    if (St45::active)
        return St45::solve(L, R);
    return solve(L, R);
}

• 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