Skip to content

Instantly share code, notes, and snippets.

@nipafx
Last active July 9, 2021 11:04
Show Gist options
  • Save nipafx/e841364700a4611ed4d7bbe2ed6be414 to your computer and use it in GitHub Desktop.
Save nipafx/e841364700a4611ed4d7bbe2ed6be414 to your computer and use it in GitHub Desktop.
Counting Red Nodes in a Colored Tree
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);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment