Storm in a Teacup - 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 clean solution to the problem, but does not pass due to the very tight time limit. See my other solution with constant optimisations that passes all tests.
#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>

using namespace std;

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

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);
    }

    // Binary Search on the answer
    long long left = 0, right = (long long)h * *max_element(value.begin(), value.end()) + 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
        function<pair<int, long long>(int, int)> dfs = [&](int node, int parent = -1) {
            // greedy, proved to be optimal
            int components = 0;
            long long leftoverSum = value[node];
            for (int child : g[node]) {
                if (child != parent) {
                    auto [c, s] = dfs(child, node);
                    components += c;
                    leftoverSum += s;
                }
            }
            if (leftoverSum >= minComponentSum) {
                ++components;
                leftoverSum = 0;
            }
            return pair{components, leftoverSum};
        };

        // if dfs from arbitrary root yields at least f components
        if (dfs(0, -1).first >= 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);
}