USACO 2022 December Contest Silver Division - Barn Tree

Problem link: here

Solution Author: Stefan Dascalescu

Problem Solution

Root the tree at node \(1\). Then do a traversal to find everyones children and find the leaf nodes. Do a reverse traversal from the leaf nodes and add orders to one of two lists based on whether there are extra haybales or a lack of haybales.

Source code

The source code in C++ can be seen below.

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

using ll = long long;

#define int ll
#define pb push_back

const int MAX_N = 2e5;

struct Order {
    int dest, numHay;
};

int h[MAX_N + 5], delta[MAX_N + 5];
bool wasVisited[MAX_N + 5];
vector<int> adj[MAX_N + 5], ansAdj[MAX_N + 5];
vector<Order> ansRoad[MAX_N + 5];
vector<int> topoOrder;
int ans;

int dfs(int node, int parent) {
    int del, delKid;

    del = delta[node];
    for (auto kid : adj[node]) {
        if (kid == parent) {
            continue;
        }

        delKid = dfs(kid, node);
        del += delKid;
        if (delKid == 0) {
            ans--;
            continue;
        }
        if (delKid > 0) {
            ansRoad[kid].push_back({node, delKid});
            ansAdj[kid].push_back(node);
        } else {
            ansRoad[node].push_back({kid, abs(delKid)});
            ansAdj[node].push_back(kid);
        }
    }
    return del;
}

void topoSort(int node) {
    wasVisited[node] = true;
    for (auto kid : ansAdj[node]) {
        if (wasVisited[kid]) {
            continue;
        }
        topoSort(kid);
    }
    topoOrder.push_back(node);
}

signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, i, tot, avg, u, v;

    cin >> n;

    tot = 0;
    for (i = 1; i <= n; i++) {
        cin >> h[i];
        tot += h[i];
    }
    avg = tot / n;
    for (i = 1; i <= n; i++) {
        delta[i] = h[i] - avg;
    }

    for (i = 0; i < n - 1; i++) {
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    ans = n - 1;
    dfs(1, 0);
    cout << ans << "\n";

    for (i = 1; i <= n; i++) {
        if (wasVisited[i]) {
            continue;
        }
        topoSort(i);
    }
    while (!topoOrder.empty()) {
        for (auto kid : ansRoad[topoOrder.back()]) {
            cout << topoOrder.back() << " " << kid.dest << " " << kid.numHay << "\n";
        }
        topoOrder.pop_back();
    }
    return 0;
}