Rainbow Roads

This problems gives us a tree where each edge has a value. It asks for the set of nodes for which we can traverse to every OTHER node in the tree without crossing two CONSECUTIVE edges with the same value.

The input size of 50k nodes lets us know off the bat that runtime will have to be linear-ish. Brute force BFS-ing from every node will be way too slow.

Imagine we have any node which has two edges which have the same value. By rule, this means that any node down either of those sub-trees CANNOT be included in our final set, since it cannot reach nodes on the other side. Therefore, the set of possible solutions is limited to nodes in any of the OTHER subtrees rooted at this node. Expanding that, the final set is the intersection of subtrees, where each subtree is rooted at some node which has two matching edges, neither of which is contained in this subtree.

The question is how to do this in linear time.

A "standard" method for intersection is when we find items we know are NOT in some set, remove them, since they can't be in the intersection. Here we apply the same principle. When we find a pair of matched edges, remove all the edges in those two subtrees from the graph, since we know they cannot be in the intersection. Whatever nodes are left in the graph after finding all matching edges must be in the intersection. The only caveat is that while removing nodes, we must check that there are no pairs of matched edges IN the subtree, in which case there are NO nodes in the intersection. If such a pair existed, it would mean any other nodes would not be able to reach some node in that subtree.

1. Iterate over all nodes not removed already
2. Find any edges emanating from that node with matching edges
3. recurse down those mathcing edge subtrees removing all nodes
4. if any path down that subtree has consecutive matched edges, there is no solution

The runtime is linear since each node is only visited once, either through the top level iteration, or because it is removed. Further, each edge is only visited at most twice. Either because we checked one of it's vertices for matching edges, or because we iterated down it during removal.

Implementation Detail

One must be very careful during implementation to not accidentally timeout. This means using BFS in the removal step to avoid deep recursion depths, as well as using hash-maps to find matching edges while checking a given vertex. Technically, using a boolean array to cache visited nodes during a removal BFS is n^2, since the array must be zeroed beforehand. In practice, though, this runs fast enough.

Since we could have to print out a huge number of nodes, buffering may be necessary. In java, a string builder should be used append before making a single syscall.

import java.awt.Point;
import java.util.*;
public class e {
static Scanner in=new Scanner(System.in);
	public static void main(String[] args) {
		int n=in.nextInt();
		ArrayList<HashMap<Integer,Integer>>al=new ArrayList<>();
		boolean[] deleted =new boolean[n];
		for(int i=0;i<n;i++)al.add(new HashMap<>());
		for(int i=0;i<n-1;i++) {
			int a=in.nextInt()-1,b=in.nextInt()-1,c=in.nextInt()-1;
			al.get(a).put(b, c);
			al.get(b).put(a, c);
		}
		for(int i=0;i<al.size();i++) if(!deleted[i]){
			boolean[] c=new boolean[n];
			HashSet<Integer>rp=new HashSet<>();
			for(int j:al.get(i).keySet()) {
				if(c[al.get(i).get(j)])rp.add(al.get(i).get(j));
				c[al.get(i).get(j)]=true;
			}
			Queue<Point>q=new LinkedList<>();
			boolean[] visited=new boolean[n];
			visited[i]=true;
			for(int j:al.get(i).keySet())if(rp.contains(al.get(i).get(j))) {//found a duplicate, elimiate subtrees
				q.offer(new Point(j,al.get(i).get(j)));
			}
			while(!q.isEmpty()) {//bfs down duplicated paths, deleting. also check if duplicate paths...in which case 0 nodes work
				Point on=q.poll();
				deleted[on.x]=true;
				for(int j:al.get(on.x).keySet()) {
					if(al.get(on.x).get(j)==on.y&&!visited[j]) {
						System.out.println(0);
						System.exit(0);
					}
					if(!deleted[j]&&j!=i) {
						q.offer(new Point(j,al.get(on.x).get(j)));
						visited[on.x]=true;
					}
				}
			}
		}
		TreeSet<Integer>ans=new TreeSet<>();
		for(int i=0;i<n;i++)if(!deleted[i])ans.add(i);
		System.out.println(ans.size());
		StringBuilder out=new StringBuilder("");
		while(!ans.isEmpty())out.append(String.format("%d\n", ans.pollFirst()+1));
		System.out.print(out);
	}
}