Simple Edit Distance

object EditDistance {
  def within(s1: String, s2: String, n: Int = 1): Boolean = {
    if (s1 != s2 && Math.abs(s1.length - s2.length) > 1) {
      return false
    }

    var s1Idx = 0
    var s2Idx = 0

    def withinN() = Math.abs(s1Idx - s2Idx) <= n

    while (Math.max(s1Idx, s2Idx) < Math.min(s1.length, s2.length) && withinN()) {
      if (s1(s1Idx) != s2(s2Idx)) {
        if (s1Idx + 1 < s1.length && s1(s1Idx + 1) == s2(s2Idx)) {
          s1Idx += 2
          s2Idx += 1
        } else {
          s1Idx += 1
          s2Idx += 2
        }
      } else {
        s1Idx += 1
        s2Idx += 1
      }
    }

    withinN()
  }
}

class WithinOneTests extends FlatSpec with Matchers {
  "Within" should "be" in {
    assert(EditDistance.within("a", "a"))
    assert(EditDistance.within("a", "ab"))
    assert(EditDistance.within("this is a fun game!", "thias is a fun game!"))
    assert(EditDistance.within("aab", "abb"))
    assert(EditDistance.within("", ""))
    assert(EditDistance.within("catrab", "atrab"))
    assert(EditDistance.within("catrab", "atxab", 2))
    assert(!EditDistance.within("aabbb", "abb"))
  }
}