Partition graph into groups of neighbours having the same class - graph

Question

Using JGraphT, I would like partition a graph into groups where each group consists of a connected sub-graph of vertices that have the same "class" (denoted using 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.

Answer 1

Sometimes you just cannot avoid writing some code
LOOP over classes
LOOP over nodes that are in class
Copy node to new graph
LOOP over edges that connect nodes in class
Copy edge to new graph
LOOP over connected components in new graph
Save component as a group graph for class

Answer 2

This is a fairly simple approach using JGraphT:
First remove edges linking neighbouring vertices that belong to different classes and then use ConnectivityInspector on the reduced graph to find connected components, which form 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();