Submission #1109748

#	Time	Username	Problem	Language	Result	Execution time	Memory
1109748		Trisanu_Das	Tree (IOI24_tree)	C++17	41 / 100	2128 ms	737656 KiB

tree

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

const int MAXN = 200005;
const long long INF = 1LL << 40;

int n, totalLeafWeight, hasMultiple;
vector<int> par, wt, adj[MAXN], subsize, rootSizes;
map<long long, long long> weightMap[MAXN];
vector<long long> subSum, prefixSum;

void initSubtree(int node) {
    if (adj[node].empty()) {
        subsize[node] = 1;
        totalLeafWeight += wt[node];
    } else {
        for (int child : adj[node]) initSubtree(child);
        if (wt[node] == 0) {
            subsize[node] = 1;
            for (int child : adj[node]) if (subsize[child] > 1) rootSizes.push_back(subsize[child]);
        } else {
            for (int child : adj[node]) subsize[node] += subsize[child];
            if (node == 0) rootSizes.push_back(subsize[0]);
        }
    }
}

void init(const vector<int> P, const vector<int> W) {
    par = P;
    wt = W;
    hasMultiple = 1;
    for (int w : wt) if (w > 1) hasMultiple = 0;
    n = par.size();
    for (int i = 0; i < n; i++) adj[i].clear();
    for (int i = 1; i < n; i++) adj[par[i]].push_back(i);
    totalLeafWeight = 0;
    subsize.assign(n, 0);
    subSum.assign(n, 0);
    rootSizes.clear();
    initSubtree(0);
    sort(rootSizes.begin(), rootSizes.end());
    prefixSum.clear();
    prefixSum.push_back(rootSizes[0]);
    for (int i = 1; i < rootSizes.size(); i++) prefixSum.push_back(prefixSum[i - 1] + rootSizes[i]);
}

long long compute(int node, int low, int high) {
    weightMap[node].clear();
    if (adj[node].empty()) {
        subSum[node] = low;
        return 1LL * wt[node] * low;
    }
    long long res = 0;
    subSum[node] = 0;
    weightMap[node][wt[node]] = INF;
    for (int child : adj[node]) {
        res += compute(child, low, high);
        subSum[node] += subSum[child];
        for (const auto& [k, v] : weightMap[child]) weightMap[node][k] += v;
    }
    if (subSum[node] > high) {
        long long excess = subSum[node] - high;
        for (auto& [curW, quota] : weightMap[node]) {
            long long use = min(excess, quota);
            res += use * curW;
            excess -= use;
            weightMap[node][curW] -= use;
            if (excess == 0) break;
        }
        subSum[node] = high;
    }
    map<long long, long long> newWeights;
    long long remaining = subSum[node] - low;
    for (const auto& [curW, quota] : weightMap[node]) {
        long long use = min(quota, remaining);
        newWeights[curW] += use;
        remaining -= use;
        if (remaining == 0) break;
    }
    swap(weightMap[node], newWeights);
    return res;
}

long long query(int low, int high) {
    if (hasMultiple) {
        long long ans = 1LL * totalLeafWeight * low;
        int sz = rootSizes.size();
        int l = 0, r = sz - 1, limit = sz;
        while (l <= r) {
            int mid = (l + r) / 2;
            if (1LL * rootSizes[mid] * low > high) {
                limit = mid;
                r = mid - 1;
            } else l = mid + 1;
        }
        if (limit < sz) {
            ans += 1LL * (prefixSum[sz - 1] - (limit ? prefixSum[limit - 1] : 0)) * low;
            ans -= 1LL * (sz - limit) * high;
        }
        return ans;
    }
    return compute(0, low, high);
}

Compilation message (stderr)

tree.cpp: In function 'void init(std::vector<int>, std::vector<int>)':
tree.cpp:44:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   44 |     for (int i = 1; i < rootSizes.size(); i++) prefixSum.push_back(prefixSum[i - 1] + rootSizes[i]);
      |                     ~~^~~~~~~~~~~~~~~~~~

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