Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- /**
- * @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);
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement