Counting Red Nodes in a Colored Tree

Tree.java

import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Stream;

// implementation of https://twitter.com/peterzllr/status/1413216826708877318

public class Tree {

	public static void main(String[] args) {
		var tree = new Red(
				new Black(
						new Red(
								new Leaf(0),
								new Leaf(1)
						),
						new Leaf(2)
				),
				new Red(
						new Leaf(3),
						new Leaf(4)
				)
		);

		System.out.println(countRed(tree));
		System.out.println(countRed_stream(tree));
	}

	static long countRed(ColoredTree tree) {
		AtomicLong count = new AtomicLong();
		tree.nodes().forEach(node -> {
			switch (node) {
				case Red red -> count.incrementAndGet();
				default -> { }
			}
		});
		return count.get();
	}

	static long countRed_stream(ColoredTree tree) {
		return tree.nodes()
				.filter(node -> node instanceof Red)
				.count();
	}

	sealed interface ColoredTree
			permits Leaf, Red, Black {

		Stream<ColoredTree> nodes();

	}

	record Leaf(int n) implements ColoredTree {

		public Stream<ColoredTree> nodes() {
			return Stream.of(this);
		}

	}

	record Red(ColoredTree left, ColoredTree right) implements ColoredTree {

		public Stream<ColoredTree> nodes() {
			// if this nested concatenation of streams turns out to be too slow,
			// feel free to implement an `Iterator<ColoredTree>` and turn that into a stream
			var children = Stream.concat(left.nodes(), right.nodes());
			return Stream.concat(Stream.of(this), children);
		}

	}

	record Black(ColoredTree left, ColoredTree right) implements ColoredTree {

		public Stream<ColoredTree> nodes() {
			var children = Stream.concat(left.nodes(), right.nodes());
			return Stream.concat(Stream.of(this), children);
		}

	}

}