CODEWITHSUNDEEP
javajava-stream
John Doe
John Doe

Reputation: 215

Recursively flatten a list with streams

I have tree-like structure with internal nodes and terminal nodes:

public interface Node
{
}

public class InternalNode implements Node {
    private List<Node> nodes;
}

public class TerminalNode implements Node {
    private String label;
}

I now have a List<Node> which I want to flatten. Here, flattening means that I want to replace an internal nodes recursively by its children, until all internal nodes are replaced by terminals.

I came up with this function:

private static List<Node> flatten(final List<Node> nodes) {
    return nodes
            .stream()
            .map(node -> {
                if (node instanceof InternalNode) {
                    return flatten(((InternalNode) node).getNodes());
                }
                return Collections.singletonList(node);
            })
            .flatMap(List::stream)
            .collect(Collectors.toList());
}

This seems to do its job. However, I am wondering if there is a better implementation possible. It seems odd that I first have to wrap a TerminalNode into a singleton list (of type List<TerminalNode>) via Collections.singletonList(node) and then later I have to convert that singleton list back to the node again via flatMap(List::stream).

Is there a way to avoid this useless Collections.singletonList(node) followed by flatMap(List::stream) for the terminal nodes?

Upvotes: 4

Views: 3822

Answers (1)

JB Nizet
JB Nizet

Reputation: 691755

You could just use flatMap directly:

private static Stream<TerminalNode> flatten(final List<Node> nodes) {
    return nodes
            .stream()
            .flatMap(node -> {
                if (node instanceof InternalNode) {
                    return flatten(((InternalNode) node).getNodes());
                }
                return Stream.of((TerminalNode) node);
            });
}

If you want a List, you can just collect the result of that method call.

Upvotes: 10

Related Questions