USACO December 2022, Silver Problem 1

link

For my solution, I ran two DFS:

For DFS 1, if a node had an excess number of haybales, meaning that the number of haybales at the node was than enough to satisfy itself and the nodes underneath it (also found with the same DFS), I would move all the excess upwards to the parent. Thus, by the end of this, there would be some parent nodes with exactly enough excess to supply to their children, and some subtrees with all nodes equal to the average value.

On DFS 2, I would simply spread the excess from the parents to the child nodes.

For each DFS iteration that had haybales involved, I would push it into a vector, and by the end, I would output the vector size and all of its contents in order.

However, my code wasn’t able to pass any of the cases besides the first, and after examination of the other test cases, it seems to be the same number of moves. I’m not sure if this has anything to do with my abstract solution, or the implementation itself. When I ran the instructions produced from my solution on the input, I would get most nodes equal to the correct value, but a few with wrong values. Any help would be appreciated.

Solution Code:

#include <bits/stdc++.h>
#define endl '\n'

using namespace std;

vector<long long> h(200005);
vector<long long> cn(200005);
vector<long long> adj[200005];
vector<pair<pair<long long, long long>, long long>> moves;
vector<bool> visited(200005);
long long avg = 0;
long long give = 0;

void dfs(long long n) { //1st DFS
    visited[n] = true;
    for(long long node : adj[n]) {
        if(visited[node]) {
            continue;
        }
        dfs(node);
        cn[n] += cn[node] + max(avg - h[node], (long long) 0); //calculating the needs of the subtree for this node
        if(give > 0) {
            moves.push_back({{node, n}, give});
            h[n] += give;
        }
    }
    if(h[n] > avg && h[n] - avg > cn[n]) {
        give = h[n] - avg - cn[n];
        h[n] -= give;
    }
    else { 
        give = 0;
    }
}

void adfs(long long n, long long give) { //2nd DFS
    visited[n] = true;
    h[n] += give;
    for(long long node : adj[n]) {
        if(visited[node]) {
            continue;
        }
        long long g = 0;
        if(h[node] < avg || cn[node] > 0) {
            if(h[node] < avg) {
                g += avg - h[node];
                if(cn[node] > 0) {
                    g+=cn[node];
                }
            }
            else if(h[node] > avg) {
                g = cn[node] - (h[node] - avg);
            }
            else {
                g = cn[node];
            }
            if(g > 0) {
                moves.push_back({{n, node}, g});
            }
            h[n] -= g;
            adfs(node, g);
        }
        else {
            adfs(node, 0);
        }
    }
    cn[n] = 0;
}

int main() {    
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    long long n; cin>>n;
    for(long long i = 1; i <= n; i++) {
        cin>>h[i];
        avg+=h[i];
    }

    avg/=n;

    for(long long i = 0; i < n-1; i++) {
        long long a, b; cin>>a>>b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    dfs(1);

    for(long long i = 1; i <= n; i++) {
        visited[i] = false;
    }

    adfs(1, 0);

    cout<<moves.size()<<endl;
    for(pair<pair<long long, long long>, long long> i : moves) {
        cout<<i.first.first<<" "<<i.first.second<<" "<<i.second<<endl;
    }

}

Perhaps try generating some small test cases and seeing if you can reproduce this incorrect behavior? (as suggested in How to ask for help on a problem)

Thanks, I realized this line:

cn[n] += cn[node] + max(avg - h[node], (long long) 0);

Should have been:

cn[n] += cn[node] + (avg - h[node]);