How to prune a graph given distance K from a node with jgrapht?

I built a graph using jgrapht API. I have a directed graph. Given a node N, I have to create a subgraph with all connected neighbours with distance K.

So basically given a node, and distance K, I have to prune the graph so that only the connected nodes with distance K remains.

I have an idea to implement it by hand.

I can generate the shortest path between all pairs from the list of nodes.
After that, I can get rid of nodes that are beyond the distance K.

However, this would result in the comparison between all the nodes and would like to know whether there is a better way to do this?

Moreover, wondering jgrapht has an API to do this already.

(I have already looked into the API of jgrapht but have not found any such API)

Solution

I assume that distance is defined as the length of the shortest path in a weighted graph, where the length of a path is given by the sum of its edge weights. I also assume that it is only required that all neighbors are within a given maxDistance from input vertex N, and that it is not required that two of those neighbors must also be within maxDistance of each other.

The simplest approach involves:

For a given input vertex N, determine all vertices that are at most maxDistance away from N.
Return an induced subgraph on N plus its (indirect) neighbors that are at most maxDistance units away.

public <V,E> Graph<V, E> getSubgraph(Graph<V,E> graph, V source, double maxDistance){
   
    //Compute a shortest. Optionally we can limit the search to vertices that are maxDistance away
    DijkstraShortestPath<V,E> ds = new DijkstraShortestPath<>(graph, maxDistance);
    ShortestPathAlgorithm.SingleSourcePaths<V, E> shortestPaths = ds.getPaths(source);
    
    //Collect all neighboring vertices that are at most maxDistance units away
    Set<V> neighbors = graph.vertexSet().stream().filter(v -> shortestPaths.getWeight(v) <= maxDistance).collect(Collectors.toSet());
    
    //Return an induced subgraph on those vertices
    return new AsSubgraph<>(graph, neighbors);
}