Partition graph into groups of neighbours having the same class

How can I use JGraphT to divide a graph into sub-graphs where each sub-graph is a connected group of vertices with the same "class" (indicated by colors below)?

Example (the desired groups are in red):

I think this is a rather simple demand and yet I can't find a (built-in) way to do this. I notice there is a PartitioningImpl class, that one constructs using a List<Set<V>> classes, yet I don't see a way to use this to partition a graph.

Ideally, I'd provide something with my graph and vertex classes (a map of V-->Integer for instance) and it would return something like a List<Set<V>> of partitioned vertex groups.

This is a fairly simple approach using JGraphT:

• Remove edges between adjacent vertices that are in distinct classes.
• Then, apply ConnectivityInspector to the resulting graph to identify the connected components, which will represent the groups.

SimpleGraph<V, E> graph; // given by user
Map<V, Integer> classes; // given by user

List<E> toRemove = new ArrayList<>();
graph.edgeSet().forEach(e -> {
    V a = graph.getEdgeSource(e);
    V b = graph.getEdgeTarget(e);
    if (!classes.get(a).equals(classes.get(b))) {
        toRemove.add(e);
    }
});

graph.removeAllEdges(toRemove);
ConnectivityInspector<V, E> ci = new ConnectivityInspector<>(graph);
List<Set<V>> groups = ci.connectedSets();