1. Compose a map that relates each MNIST digit to a list of (score, feature) pairs:

Map<Integer,List<Pair<Double,INDArray>>> listsByDigit = new HashMap<>();

2. Iterate through each and every test feature, calculate the reconstruction error, make a score-feature pair for the purpose of displaying the sample with a low reconstruction error:

for( int i=0; i<featuresTest.size(); i++ ){
 INDArray testData = featuresTest.get(i);
 INDArray labels = labelsTest.get(i);
 for( int j=0; j<testData.rows(); j++){
 INDArray example = testData.getRow(j, true);
 int digit = (int)labels.getDouble(j);
 double score = net.score(new DataSet(example,example));
 // Add (score, example) pair to the appropriate list
 List digitAllPairs = listsByDigit.get(digit);
 digitAllPairs.add(new Pair<>(score, example));
 }
 }

3. Create a custom comparator to sort the map:

Comparator<Pair<Double, INDArray>> sortComparator = new Comparator<Pair<Double, INDArray>>() {
 @Override
 public int compare(Pair<Double, INDArray> o1, Pair<Double, INDArray> o2) {
 return Double.compare(o1.getLeft(),o2.getLeft());
 }
 };

4. Sort the map using Collections.sort():

for(List<Pair<Double, INDArray>> digitAllPairs : listsByDigit.values()){
 Collections.sort(digitAllPairs, sortComparator);
 }

5. Collect the best/worst data to display in a JFrame window for visualization:

List<INDArray> best = new ArrayList<>(50);
 List<INDArray> worst = new ArrayList<>(50);
 for( int i=0; i<10; i++ ){
 List<Pair<Double,INDArray>> list = listsByDigit.get(i);
 for( int j=0; j<5; j++ ){
 best.add(list.get(j).getRight());
 worst.add(list.get(list.size()-j-1).getRight());
 }
 }

6. Use a custom JFrame implementation for visualization, such as MNISTVisualizer, to visualize the results:

//Visualize the best and worst digits
 MNISTVisualizer bestVisualizer = new MNISTVisualizer(imageScale,best,"Best (Low Rec. Error)");
 bestVisualizer.visualize();
 MNISTVisualizer worstVisualizer = new MNISTVisualizer(imageScale,worst,"Worst (High Rec. Error)");
 worstVisualizer.visualize();