Replace Set by Map for later extension

2 files changed, 13 insertions, 8 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,
diff --git a/spellchecker/test/SpellCorrectorTest.java b/spellchecker/test/SpellCorrectorTest.java
index a00e236..d658914 100644
--- a/spellchecker/test/SpellCorrectorTest.java
+++ b/spellchecker/test/SpellCorrectorTest.java
@@ -12,18 +12,20 @@ import org.junit.Before;
 public class SpellCorrectorTest {
 
     private CorpusReader fullCR;
+    private ConfusionMatrixReader cmr;
 
     @Before
     public void setUp() throws IOException {
         fullCR = new MockCorpusReader();
+        cmr = new ConfusionMatrixReader();
     }
 
     private void checkGetCandidateWords(CorpusReader cr, String word,
             Set<String> expResult)
             throws IOException {
         System.out.println("getCandidateWords(" + word + ")");
-        SpellCorrector instance = new SpellCorrector(cr, null);
-        HashSet<String> result = instance.getCandidateWords(word);
+        SpellCorrector instance = new SpellCorrector(cr, cmr);
+        Set<String> result = instance.getCandidateWords(word).keySet();
         // quick check: are the results of the same size?
         assertEquals(expResult.size(), result.size());
         // verbose test: are all letters as expected?