Tea - BIO 2022 Round 2

#include <bits/stdc++.h>

using namespace std;

int main() {
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
    int h, t; cin >> h >> t;
    if (h == 2) { // just cover the case where all nodes are leaves separately
        int x; cin >> x >> x; // input the edge between 1 and 2
        long long sum[2] {};
        while (t--) {
            int a, b;
            long long tea;
            cin >> a >> b >> tea;
            if (a != b) sum[0] += tea, sum[1] += tea;
            else if (a == 1) sum[0] += tea;
            else sum[1] += tea;
        }
        cout << sum[0] << endl << sum[1] << endl;
        return 0;
    }
    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);
    }

    vector<int> pos(h); // pos -1 means leaf, otherwise this is the position of the node in the non-leaf path
    for (int i = 0; i < h; ++i) {
        if (g[i].size() == 1) pos[i] = -1;
    }

    // Now order the non-leaf nodes
    int startNode = -1;
    for (int i = 0; i < h; ++i) {
        if (pos[i] == -1) continue;
        int nonLeafDegree = 0;
        for (int j : g[i]) {
            if (pos[j] != -1) ++nonLeafDegree;
        }
        if (nonLeafDegree < 2) {
            startNode = i;
            break;
        }
    }
    assert(startNode != -1);

    vector<int> order;
    int cur = startNode, parent = -1;
    while (cur != -1) {
        pos[cur] = order.size();
        order.push_back(cur);
        int nxt = -1;
        for (int j : g[cur]) {
            if (pos[j] != -1 && j != parent) {
                nxt = j;
                break;
            }
        }
        parent = cur;
        cur = nxt;
    }

    // Now use prefix sums to store values over this path
    int n = order.size();
    vector<long long> prefixPath(n);
    vector<long long> leafTea(h); // Pun not intended - total tea in each leaf

    while (t--) {
        int a, b;
        long long tea;
        cin >> a >> b >> tea;
        --a, --b;

        if (pos[a] == -1) leafTea[a] += tea, a = g[a][0];
        if (pos[b] == -1) leafTea[b] += tea, b = g[b][0];

        if (pos[a] > pos[b]) swap(a, b);
        prefixPath[pos[a]] += tea;
        if (pos[b]+1 < n) prefixPath[pos[b]+1] -= tea;
    }

    for (int i = 1; i < n; ++i) prefixPath[i] += prefixPath[i-1];
    for (int i = 0; i < h; ++i) {
        if (pos[i] == -1) cout << leafTea[i] << endl;
        else cout << prefixPath[pos[i]] << endl;
    }
    return 0;
}