Replace Set by Map for later extension

1 files changed, 9 insertions, 6 deletions

diff --git a/spellchecker/src/SpellCorrector.java b/spellchecker/src/SpellCorrector.java
index 37020fb..6c9225b 100644
--- a/spellchecker/src/SpellCorrector.java
+++ b/spellchecker/src/SpellCorrector.java
@@ -1,5 +1,6 @@
 
-import java.util.HashSet;
+import java.util.HashMap;
+import java.util.Map;
 
 public class SpellCorrector {
 
@@ -40,10 +41,10 @@ public class SpellCorrector {
      * substitution or transposition.
      *
      * @param word The (wrong) word to find (corrected) candidates for.
-     * @return A set of words with candidate words.
+     * @return Candidate words mapping to the noisy channel model probability.
      */
-    public HashSet<String> getCandidateWords(String word) {
-        HashSet<String> ListOfWords = new HashSet<>();
+    public Map<String, Double> getCandidateWords(String word) {
+        Map<String, Double> candidates = new HashMap<>();
 
         // tries to add word2 to the list of words. This word2 was generated
         // from "word" by changing "error" to "correct".
@@ -52,7 +53,9 @@ public class SpellCorrector {
                 return;
             }
 
-            ListOfWords.add(word2);
+            double p = candidates.getOrDefault(word2, 0.0);
+
+            candidates.put(word2, p);
         };
 
         makeWordInsertion(word, push);
@@ -60,7 +63,7 @@ public class SpellCorrector {
         makeWordDeletion(word, push);
         makeWordTransposition(word, push);
 
-        return ListOfWords;
+        return candidates;
     }
 
     private void makeWordInsertion(String word,