CNSTREE2 - Editorial

editorial
, , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093

DIFFICULTY:

Easy - Medium

PREREQUISITES:

Trees

PROBLEM:

You are given a tree, each vertex has a value in \{-1, 0, 1\}.
The tree is called consistent if there exists a way to assign the labels 1 to N to the vertices such that, for each vertex u:

  • u has value -1 iff its label is smaller than all its neighbors’ labels.
  • u has value 1 iff its label is larger than all its neighbors’ labels.
  • u has value 0 iff neither of the above.

Given a tree, decide if it is consistent; and if it is, find a valid labeling.

EXPLANATION:

Let’s decide when a tree is consistent.

There are some obvious initial conditions.

  • If two vertices with value -1 are adjacent, then the tree cannot be consistent; since one of them has to end up with a larger value than the other.
  • Similarly, two vertices with value 1 cannot be adjacent to each other.
  • A vertex with value 0 must have at least two neighbors; since it needs both a smaller and a larger neighbor.

These conditions are necessary, but not sufficient just yet.
A simple example here is a line with three vertices, where the leaves both have value -1 and the middle has value 0 - it satisfies the above conditions, but no valid labeling exists because each leaf must be smaller than the middle, but then the middle cannot be smaller than any of its neighbors.

However, we wouldn’t run into this problem if one leaf had -1 and one had 1 instead.

Generalizing this idea, we obtain the following conclusion:

  • Consider any connected set of vertices that all have value 0.
    This component must be adjacent to at least one -1, and also adjacent to at least one 1.

This condition, along with the previous three, is now sufficient for the tree to be consistent!

The proof of this will be via construction.

Suppose all the four conditions above hold.
Let there be X vertices with value -1 and Y vertices with value 1, in total.

We first assign 1, 2, \ldots, X to the vertices with value -1, in any order.
Then we assign N, N-1, \ldots, N-Y+1 to the vertices with value 1, again in any order.
This handles them all; now we only need to worry about vertices with value 0.

Consider some connected component of zeros.
This component is itself a tree, so there are two possibilities: either it’s a singleton vertex, or it’s a tree of size \ge 2.

If it’s a singleton vertex, this vertex must be adjacent to both a 1 and a -1 by assumption.
This means we can assign any free value to it and be fine; since our free values are in [X+1, N-Y].

Finally, we look at the component if it has \ge 2 vertices.
If this is the case, the component must have a ‘leaf’ vertex, i.e. a 0-value vertex that’s adjacent to only one other 0-value vertex.
Pick any such ‘leaf’, say v.
Then,

  • v must be adjacent to either a 1 or a -1.
  • If it’s adjacent to a 1 (but not a -1), give it the highest remaining value; and treat it as a 1-vertex in the future.
    • This ensures its neighbor receives a smaller value than it; hence satisfying it (the 1 its neighboring already has a larger value.)
    • Note that the remaining component of 0’s is still adjacent to both a 1 and a -1.
  • If it’s adjacent to a -1 (but not a 1), give it the smallest remaining value; and treat it as a -1-vertex in the future.
    • This ensures its neighbor receives a larger value than it; hence satisfying it.
    • The remaining component of 0’s is still adjacent to both a 1 and a -1.
  • If it’s adjacent to both a 1 and a -1, ignore this leaf for now, and pick a different one.

Eventually, we’ll reach a situation where all remaining ‘leaves’ of the component have both 1 and -1 adjacent to them.

Now we can do the following:

  • Pick some vertex in the component that’s adjacent to a -1. Call it s.
  • Root the component at s.
  • Give s the smallest remaining free value, then the children of s the next smallest free values, and so on.
    Essentially, fill values small to large, from root downwards along the component.

This works because:

  • s is satisfied since the -1 adjacent to it has a smaller value, and its children have larger values.
  • All ‘internal’ vertices are satisfied, since their parents are smaller and children are larger.
  • All leaves are satisfied because we ended up in a state where leaves were adjacent to both 1 and -1 anyway.

This can all be implemented in linear time.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Tester's code (C++) 
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18

mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());

void Solve() 
{
    int n; cin >> n;
    
    vector <int> a(n + 1);
    for (int i = 1; i <= n; i++){
        cin >> a[i];
    }
    
    vector<vector<int>> adj(n + 1);
    for (int i = 1; i < n; i++){
        int u, v; cin >> u >> v;
        
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    vector<vector<int>> g(n + 1);
    vector <int> indeg(n + 1);
    
    auto add = [&](int u, int v){
        g[u].push_back(v);
        indeg[v]++;
    };
    
    for (int i = 1; i <= n; i++){
        for (int v : adj[i]) if (v > i && ((a[i] != 0) || (a[v] != 0))){
            if (a[i] == -1 || a[v] == 1) add(i, v);
            else add(v, i);
        }
    }
    
    vector <int> pos(n + 1, 0), neg(n + 1, 0);
    for (int i = 1; i <= n; i++){
        if (a[i] == 0){
            for (int v : adj[i]){
                if (a[v] == -1){
                    neg[i]++;
                } else if (a[v] == 1){
                    pos[i]++;
                }
            }
        }
    }
    
    vector <bool> vis(n + 1, false);
    vector <int> pp(n + 1);
    
    for (int i = 1; i <= n; i++) if (a[i] == 0 && !vis[i] && neg[i]){
        queue <int> q;
        q.push(i);
        vis[i] = true;
        
        vector <int> ord;
        
        while (!q.empty()){
            int u = q.front(); q.pop();
            
            ord.push_back(u);
            
            for (int v : adj[u]) if (!vis[v] && a[v] == 0){
                vis[v] = true;
                pp[v] = u;
                q.push(v);
            }
        }
        
        for (int i = (int)ord.size() - 1; i >= 0; i--){
            int u = ord[i];
            if (i){
                pos[pp[u]] |= pos[u];
                if (pos[u]) add(pp[u], u);
                else add(u, pp[u]);
            }
        }
    }
    
    vector <int> ans(n + 1);
    int curr = 1;
    queue <int> q;
    for (int i = 1; i <= n; i++) if (indeg[i] == 0) q.push(i);
    
    while (!q.empty()){
        int u = q.front(); q.pop();
        ans[u] = curr++;
        
        for (int v : g[u]){
            indeg[v]--;
            if (indeg[v] == 0){
                q.push(v);
            }
        }
    }
    
    if (curr <= n){
        cout << -1 << "\n";
        return;
    }
    
    for (int i = 1; i <= n; i++){
        int ge = 0, le = 0;
        for (int v : adj[i]){
            ge += ans[v] > ans[i];
            le += ans[v] < ans[i];
        }
        
        int sg = (le > 0) - (ge > 0);
        if (sg != a[i]){
            cout << -1 << "\n";
            return;
        }
    }
    
    for (int i = 1; i <= n; i++){
        cout << ans[i] << " \n"[i == n];
    }
}

int32_t main() 
{
    auto begin = std::chrono::high_resolution_clock::now();
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    // freopen("in",  "r", stdin);
    // freopen("out", "w", stdout);
    
    cin >> t;
    for(int i = 1; i <= t; i++) 
    {
        //cout << "Case #" << i << ": ";
        Solve();
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
    cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n"; 
    return 0;
}
CNSTREE - Editorial