diff --git a/src/main/java/p2/wordsuggestor/NGramToNextChoicesMap.java b/src/main/java/p2/wordsuggestor/NGramToNextChoicesMap.java index 1e1a6c2dd70f935bd401e13618648dcabcd7f090..7f66f8433375cce24d2c4edd2a9b6f9502bee5da 100644 --- a/src/main/java/p2/wordsuggestor/NGramToNextChoicesMap.java +++ b/src/main/java/p2/wordsuggestor/NGramToNextChoicesMap.java @@ -73,10 +73,16 @@ public class NGramToNextChoicesMap { Comparator<Item<String, Integer>> comp = new LargeValueFirstItemComparator<String, Integer>(); if (k < 0) { - InsertionSort.sort(afterNGrams, comp); + QuickSort.sort(afterNGrams, comp); } else { - // You must fix this line toward the end of the project - throw new NotYetImplementedException(); + Comparator<Item<String, Integer>> revComp = (o1, o2) -> comp.compare(o2, o1); + TopKSort.sort(afterNGrams, k, revComp); + int cutoff = Math.min(k, afterNGrams.length); + for (int i = 0; i < cutoff / 2; i++) { + Item<String, Integer> temp = afterNGrams[i]; + afterNGrams[i] = afterNGrams[cutoff - i - 1]; + afterNGrams[cutoff - i - 1] = temp; + } } String[] nextWords = new String[k < 0 ? afterNGrams.length : k];