Storm in a Teacup (faster) - BIO 2024 Round 2

/**
 * @file teacup.cpp
 * @version 1.0
 * @date 2024-03-30
 *
 * usage:
 *      Read from / write to default input.txt and output.txt
 *          teacup.exe
 *      Read from / write to custom files:
 *          teacup.exe in.txt out.txt
 */
// This is a constant-optimised solution - the problem had a very tight time limit. My other solution is cleaner and less complex, although it has a larger constant.
#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
#include <numeric>

#include <cassert>

using namespace std;

void fileIO(int argc, char *argv[]);

const int MX = 1<<20;
int parent[1+MX], components[1+MX];
long long leftoverSum[1+MX];

int main(int argc, char *argv[]) {
    fileIO(argc, argv); // remove for standard input/output
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    int h, f; cin >> h >> f;
    vector<int> value(h);
    for (int &x : value) cin >> x;

    vector<vector<int>> g(h);
    for (int i = 0; i < h-1; ++i) {
        int u, v; cin >> u >> v;
        --u, --v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    // precalculate order using bfs to topo sort, do not need queues
    vector<int> order;
    order.push_back(0);
    parent[0] = h; // no parent
    for (size_t nxt = 0; nxt < order.size(); ++nxt) {
        int node = order[nxt];
        for (int child : g[node]) {
            if (child != parent[node]) {
                order.push_back(child);
                parent[child] = node;
            }
        }
    }
    assert((int)order.size() == h);

    // reindex nodes in this order
    vector<vector<int>> g2(h);
    vector<int> iAt(h);
    for (int i = 0; i < h; ++i) iAt[order[i]] = i;
    assert(iAt[0] == 0);
    for (int i = 0; i < h; ++i) {
        for (int j : g[i]) {
            if (j == parent[i]) continue;
            g2[iAt[i]].push_back(iAt[j]);
            assert(iAt[i] < iAt[j]);
        }
    }
    vector<int> value2(h);
    for (int i = 0; i < h; ++i) value2[iAt[i]] = value[i];

    // Binary Search on the answer
    long long sum = accumulate(value.begin(), value.end(), 0LL);
    long long left = 0, right = sum/f + 1;
    while (left+1 < right) {
        long long minComponentSum = (left + right) / 2;
        /* check if the graph can be split into f components,
        each with sum >= minComponentSum */

        // dp on subtrees

        bool works = false;

        // not using node : order and child : g[node] since it is slower
        for (int node = (int)order.size()-1; node >= 0; --node) {
            components[node] = 0, leftoverSum[node] = value2[node];
            for (size_t j = 0; j < g2[node].size(); ++j) {
                int child = g2[node][j];
                // can avoid making sure child is not parent node, since parent's values are 0
                components[node] += components[child];
                leftoverSum[node] += leftoverSum[child];
            }
            if (leftoverSum[node] >= minComponentSum) {
                leftoverSum[node] = 0;
                ++components[node];
            }
            // if (components[node] >= f) {
            //     works = true;
            //     break;
            // }
        }

        // if dfs from arbitrary root yields at least f components
        if (works || components[0] >= f) left = minComponentSum;
        else right = minComponentSum;
    }
    cout << left << '\n';
    return 0;
}

void fileIO(int argc, char *argv[]) {
    const char * in = "input.txt", * out = "output.txt";
    if (argc > 1) in = argv[1];
    if (argc > 2) out = argv[2];
    freopen(in, "r", stdin);
    freopen(out, "w", stdout);
}