Balanced Flow Tree#12237

Balanced Flow Tree (Editorial)

Problem Author:	Imran AKKI
Prepared by:	Imran AKKI
Edited by:	Hasan Saleh (MathModel)

Solution a single-pass leaf-peeling greedy: keep $rem [u] = b [u]$ \text{rem}[u]=b[u] (the remaining out-in each vertex must obtain), push all leaves into a queue, and repeatedly remove any leaf $u$ with its unique unused edge $e = (u, p)$ e=(u,p). When $u$ has only that one edge left the contribution of $e$ to $u$ must be $+ 1$ (if $u \to p$ u \rightarrow p) or $- 1$ (if $p \to u$ p \rightarrow u), so $rem [u]$ must be $+ 1$ or $- 1$ ; orient edge $e$ accordingly, mark it used, remove $u$ (decrement degrees), and propagate the effect to the neighbor by $rem [p] = t e x t r e m [p] + rem [u]$ \text{rem}[p]=text{rem}[p]+\text{rem}[u].

If at any point a leaf $u$ has $rem [u]$ not in ${+ 1, - 1}$ report NO.

Continue until all edges are oriented (then YES) or you detect a contradiction. This is correct because any valid global orientation induces the same forced choice whenever a vertex becomes a leaf (an edge to a leaf can only give $\pm 1$ ), so our greedy matches any solution and, by induction on the number of vertices, reconstructs it if one exists; conversely, if the algorithm fails at some leaf no global orientation can satisfy that leaf's required remaining value.

Total complexity is $O (n)$ .

// Author: Imran AKKI (imranakki)
#include <bits/stdc++.h>
#define ll long long
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define lp(i, a, n) for (ll i = a; i < n; i++)
#define MAXN 1000005
#define MOD (ll)(1e9 + 7)
using namespace std;

void solve()
{
    ll n; cin >> n;
    vector<ll> b(n);
    for(auto &e:b) cin >> e;

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

    vector<ll> deg(n, 0);
    queue<ll> q;
    for(ll i = 0; i < n; i++) {
        deg[i] = adj[i].size();
        if(deg[i] == 1) q.push(i);
    }

    vector<pair<ll, ll>> ans;
    vector<bool> vis(n, 0);
    while(!q.empty()) {
        ll u = q.front(); q.pop();
        if(vis[u] || deg[u] != 1) continue;

        if(b[u] != 1 && b[u] != -1) {
            cout << "NO\n";
            return;
        }

        ll p = -1;
        for(auto v: adj[u]) {
            if(!vis[v]){
                p = v;
                break;
            }
        }

        bool u_to_p = (b[u] == 1);
        if(u_to_p){
            ans.push_back({u, p});
            b[p]++;
        }else{
            ans.push_back({p, u});
            b[p]--;
        }

        vis[u] = 1;
        deg[p]--;
        if(deg[p] == 1) {
            q.push(p);
        }
        deg[u] = 0;
    }
    if((ll)ans.size() != n - 1) {
        cout << "NO\n";
    }else{
        cout << "YES\n";
        for(auto &e: ans) cout << e.first + 1 << " " << e.second + 1 << "\n";
    }
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    solve();
    return 0;
}

C++

Complexity : $O (n)$