Simple spell checker based on dynamic programming

spellchecker.scala

abstract class Spelling
case class CorrectSpelling(word: String) extends Spelling
case class IncorrectSpelling(word: String, suggestions: List[String]) extends Spelling

case class Spellchecker(dictionary: String) {
  private val numSuggestions = 10
  private val maxDistance = 5
  // set of valid words (replace with trie for space efficiency)
  private val set = readDict(dictionary)
  private val heap = new java.util.PriorityQueue[(Int, String)](
    numSuggestions, 
    new java.util.Comparator[(Int, String)] {
      override def compare(a: (Int, String), b: (Int, String)): Int = {
        if (a._1 > b._1) return -1
        if (a._1 < b._1) return 1
        return 0
      }
    }
  )

  private def readDict(path: String): java.util.HashSet[String] = {
    val iter = scala.io.Source.fromFile(path).getLines
    val set = new java.util.HashSet[String]()
    while (iter.hasNext) {
      set.add(iter.next.toLowerCase)
    }
    return set
  }

  private def balance(heap: java.util.PriorityQueue[(Int, String)], size: Int): Unit = {
    while (heap.size > size) {
      heap.poll()
    }
  }

  private def distance(s: String, t: String): Int = {
    val n = s.length + 1
    val m = t.length + 1
    val dp = new Array[Array[Int]](n)
    for (i <- 0 until n) dp(i) = new Array[Int](m)
    for (i <- 0 until n) dp(i)(0) = i
    for (j <- 0 until m) dp(0)(j) = j

    for (i <- 1 until n) {
      for (j <- 1 until m) {
        val a = s.charAt(i-1)
        val b = t.charAt(j-1)
        dp(i)(j) = dp(i-1)(j-1) + (if (a == b) 0 else 1)
        dp(i)(j) = Math.min(dp(i)(j), dp(i)(j-1) + 1)
        dp(i)(j) = Math.min(dp(i)(j), dp(i-1)(j) + 1)
      }
    }
    return dp(n-1)(m-1)
  }

  def check(word: String): Spelling = {
    val target = word.toLowerCase
    if (set.contains(target)) return CorrectSpelling(target)
    heap.clear()
    val iter = set.iterator()
    while (iter.hasNext) {
      val s = iter.next
      val edits = distance(s, target)
      if (edits <= maxDistance) {
        heap.add((edits, s))
        balance(heap, numSuggestions)
      }
    }

    // extract suggestions and build result
    val suggestions = new Array[String](heap.size)
    var i = heap.size - 1
    while (!heap.isEmpty) {
      suggestions(i) = heap.poll()._2
      i -= 1
    }
    return IncorrectSpelling(target, suggestions.toList)
  }
}

val a = Spellchecker("words.txt")
a.check("deleete")
// res5: Spelling = IncorrectSpelling(deleete,List(delete, delate, desuete, dolente, delegate, deleted, decrete, 
//    melete, celeste, detente))
a.check("deleteed")
// res11: Spelling = IncorrectSpelling(deleteed,List(deleted, depleted, delete, deletes, deleter, delated, deleaded, 
//    deleaved, deedeed, deletery))

We can optimize suggestions in several ways, e.g. short-circuiting some of the words and using global dp matrix for all solutions.