Ball-Stackable
Contest: Codeforces Round 902 (Div. 1, based on COMPFEST 15 - Final Round)
Difficulty rating: 3300
Tags: constructive algorithms, data structures, dp, trees

With a problem title like that, there is no way this is going to be a graph problem.

Chaneka has a graph with $$$n$$$ vertices and $$$n-1$$$ edges. Some of the edges are directed and some of the edges are undirected. Edge $$$i$$$ connects vertex $$$u_i$$$ to vertex $$$v_i$$$. If $$$t_i=0$$$, edge $$$i$$$ is undirected. If $$$t_i=1$$$, edge $$$i$$$ is directed in the direction from $$$u_i$$$ to $$$v_i$$$. It is known that if you make all edges undirected, the graph becomes a tree$$$^\dagger$$$.

Chaneka wants to direct all undirected edges and colour each edge (different edges can have the same colour).

After doing that, suppose Chaneka starts a walk from an arbitrary vertex $$$x$$$ to an arbitrary vertex $$$y$$$ (it is possible that $$$x=y$$$) going through one or more edges. She is allowed to go through each edge either following the direction or opposite to the direction of the edge. She is also allowed to visit a vertex or an edge more than once. During the walk, Chaneka maintains a stack of balls that is initially empty before the walk. Each time Chaneka goes through an edge, she does the following:

- If Chaneka goes through it in the right direction, she puts a new ball with a colour that is the same as the edge's colour to the top of the stack.
- If Chaneka goes through it in the opposite direction, she removes the ball that is on the top of the stack.

A walk is stackable if and only if the stack is not empty before each time Chaneka goes through an edge in the opposite direction.

A walk is ball-stackable if and only if it is stackable and each time Chaneka goes through an edge in the opposite direction, the colour of the ball removed from the stack is the same as the colour of the edge traversed.

Is it possible to direct all undirected edges and colour each edge such that all stackable walks are also ball-stackable? If it is possible, find a construction example that uses the maximum number of different colours among all valid ways of directing and colouring. If there are multiple such solutions, output any of them.

$$$^\dagger$$$ A tree is a connected graph with no cycles.

Input format:
The first line contains a single integer $$$n$$$ ($$$2\leq n\leq10^5$$$) — the number of vertices in the graph.

The $$$i$$$-th of the next $$$n-1$$$ lines contains three integers $$$u_i$$$, $$$v_i$$$, and $$$t_i$$$ ($$$1 \leq u_i,v_i \leq n$$$; $$$0\leq t_i\leq1$$$) — an undirected edge connecting vectices $$$u_i$$$ and $$$v_i$$$ if $$$t_i=0$$$, or a directed edge from vertex $$$u_i$$$ to vertex $$$v_i$$$ if $$$t_i=1$$$. If you make all edges undirected, the graph becomes a tree.

Output format:
A single line containing $$$-1$$$ if it is impossible.

Otherwise, the output consists of $$$n$$$ lines describing your construction. The first line contains an integer $$$z$$$ representing the number of different colours used. The $$$i$$$-th of the next $$$n-1$$$ lines contains three integers $$$p$$$, $$$q$$$, and $$$c$$$ ($$$1\leq p,q\leq n$$$; $$$1\leq c\leq z$$$) — the edge connecting vertices $$$p$$$ and $$$q$$$ in the graph is directed from vertex $$$p$$$ to vertex $$$q$$$ and is coloured with colour $$$c$$$. If there are multiple such solutions, output any of them.

Note that since there should be $$$z$$$ different colours in your construction, that means each colour from $$$1$$$ to $$$z$$$ must appear at least once in the graph.

Constraints:
Time limit: 2.0 s
Memory limit: 512.0 MB

Note:
The following is the given graph.

Chaneka can direct all undirected edges and colour each edge as follows.

As an example, consider a stackable walk $$$3→1→5→2→5→4→5$$$. Let's show that that walk is ball-stackable.

1. Chaneka starts in vertex $$$3$$$. The stack is $$$[]$$$.
2. Chaneka moves to vertex $$$1$$$. She puts a ball of colour $$$3$$$. The stack is $$$[3]$$$.
3. Chaneka moves to vertex $$$5$$$. She puts a ball of colour $$$2$$$. The stack is $$$[3,2]$$$.
4. Chaneka moves to vertex $$$2$$$. She removes a ball of colour $$$2$$$ (same colour as the edge). The stack is $$$[3]$$$.
5. Chaneka moves to vertex $$$5$$$. She puts a ball of colour $$$2$$$. The stack is $$$[3,2]$$$.
6. Chaneka moves to vertex $$$4$$$. She puts a ball of colour $$$1$$$. The stack is $$$[3,2,1]$$$.
7. Chaneka moves to vertex $$$5$$$. She removes a ball of colour $$$1$$$ (same colour as the edge). The stack is $$$[3,2]$$$.

Since every time Chaneka removes a ball from the stack, it has the same colour as the edge traversed, then the walk above is ball-stackable. It can be shown that if we direct and colour the edges as shown above, any possible stackable walk is also ball-stackable.

Example input:
5
2 5 1
1 3 0
5 4 0
1 5 1

Example output:
3
1 5 2
5 4 1
2 5 2
3 1 3

Editorial:
Let's try to solve the problem if all edges are directed.
First, assign some arbitrary vertex $$$x$$$ as the root of the tree. Consider all walks that start from $$$x$$$. Notice that for a fixed vertex $$$y$$$, no matter how Chaneka walks from $$$x$$$, when she reaches $$$y$$$, the number of balls in the stack will always be the same. More precisely, if in the shortest path from $$$x$$$ to $$$y$$$ there are $$$c_1$$$ edges in the right direction and $$$c_2$$$ edges in the opposite direction, then the size of the stack in $$$y$$$ is $$$c_1-c_2$$$.
For some root $$$x$$$, let's calculate the stack size of every vertex using DFS. What happens if the stack size that is calculated is negative? This conflicts with our logic, but let's ignore them for now.
Next, find a vertex with the minimum stack size (including negative stack sizes). Then, assign this a the new root $$$x$$$ and do the same process to recalculate every stack size. Because of the way we choose the new $$$x$$$, the newly recalculated stack sizes will always be non-negative.
Consider a path from this root to some other vertex and go through that path while maintaining a stack of colourful balls. Each time we go through an edge in the right direction, we add a new ball to the stack, so we have an opportunity to use a new colour for the edge (and simultaneously for the new ball). Each time we go through an edge in the opposite direction, we remove the ball at the top of the stack, so we must colour the edge to be equal to some ball we added before (the colour of some previous edge).
We can simulate every path from this root to every other vertex by doing DFS once while maintaining a stack. Using the rules above, we get the rules for colouring the edges of the graph. In fact, simulating every path from this single root already solves the problem with all directed edges, because there are no other relationships between two edges that must have the same colour that we have not considered. Remember that we do not have to worry about removing from an empty stack because of how we smartly choose the root.
Notice that using the rules above, each edge that is directed away from the root contributes to add one new different colour to the colouring while each edge that is directed towards the root does not add any new colour. This means we just want to maximize the total number of edges that is directed away from the root.
We can expand this idea for the original problem with undirected edges. We need to find a root, then direct all undirected edges away from that root such that the number of edges directed away from that root is as many as possible. This is equivalent to finding a vertex such that the number of directed edges (ignoring all undirected edges) that is directed away from that vertex is maximum. This vertex can be found using a simple DFS. Furthermore, choosing such a vertex will automatically eliminates negative stack sizes if we set it as the root.
Once we find the optimal root and direct the undirected edges, we just do a simulation DFS like above to determine the colouring.
Time complexity: $$$O(n)$$$