Given an array of integers, return indices of the two numbers such that they add up to a specific target with Scala

Viewed 693

I practice Scala by solving algo questions and stuck on this problem from leetcode:

Given an array of integers, return indices of the two numbers such that they add up to a specific target. You may assume that each input would have exactly one solution, and you may not use the same element twice.

Example:

Given nums = [2, 7, 11, 15], target = 9,

Because nums[0] + nums[1] = 2 + 7 = 9,
return [0, 1].

I came up with the following code

def twoSum(nums: Array[Int], target: Int): Array[Int] = {
    val map = scala.collection.mutable.Map.empty[Int, Int]
    nums.zipWithIndex.foreach{ case(num, idx) => 
        map.get(idx) match {
            case Some(r) => return Array(r, idx)
            case None => map += (idx -> (target-num))
       }
    }
    Array(-1, -1)
}

It seems that the I'm adding values to map in a wrong way, as for every input I get an Array(-1,-1) as a result.

The expected soultion.

def twoSum(nums: Array[Int], target: Int): Array[Int] = {
    var map = Map.empty[Int, Int]
    for ((v, i) <- nums.zipWithIndex) {
       map.get(v) match {
           case Some(r) => return Array(r, i)
           case None => map ++= Map((target - v) -> i)
       }
    }
    Array(-1, -1)
}

Can you point on my mistake and explain how to update hashMap properly? While updating the elements with mutable Map I was referencing this post.

2 Answers

In the second solution you have a map of value to index, which is correct. But in the first one you have the opposite: a map from index to value, which is wrong. So you have to reverse the entries you add to the Map:

def twoSum(nums: Array[Int], target: Int): Array[Int] = {
  val map = scala.collection.mutable.Map.empty[Int, Int]
  nums.zipWithIndex.foreach{ case(num, idx) =>
    map.get(num) match {                     // Changed: retrieving the index of the value 
      case Some(r) => return Array(r, idx)
      case None => map += ((target-num) -> idx)  // Changed: adding (value -> index) pairs
    }
  }
  Array(-1, -1)
}

For the record, if you want to avoid using mutability and return, the simplest way to do it would be to precompute the Map in one pass, and then do the second pass to find both indexes:

import scala.collection.immutable.HashMap

def twoSum(nums: Array[Int], target: Int): Array[Int] = {
  val indexByValue = HashMap(nums.zipWithIndex: _*)
  val result = nums.indices.collectFirst(Function.unlift { i =>
    indexByValue.get(target - nums(i)).filter(_ != i).map(Array(i, _))
  })
  result.getOrElse(Array(-1, -1))
}

This is still slower than the mutable version, because it does two passes over nums and eagerly builds the whole Map, while the original mutable code only does a single pass and builds the Map lazily. To compute the Map lazily in a single pass in a functional way you can use scanLeft over an iterator or some lazy collection:

import scala.collection.immutable.HashMap

def twoSum(nums: Array[Int], target: Int): Array[Int] = {
  nums.indices.iterator
    .scanLeft(HashMap.empty[Int, Int]) { (map, i) =>
      map + (nums(i) -> i)
    }
    .zipWithIndex
    .collectFirst(Function.unlift { case (indexByValue, i) =>
      indexByValue.get(target - nums(i)).filter(_ != i).map(Array(i, _))
    })
    .getOrElse(Array(-1, -1))
}

You can emulate what you would do on an imperative language, a nested for.
(note: I changed the signature a bit to be more idiomatic, you can just call this function on the original one, adapting the inputs and outputs)

import scala.collection.immutable.ArraySeq // Immutable array, only exists on 2.13

def twoSum(nums: ArraySeq[Int], target: Int): Option[(Int, Int)] =
  Iterator
    .range(start = 0, end = nums.length)
    .map { i =>
      Iterator
        .range(start = (i + 1), end = nums.length)
        .collectFirst {
          case j if ((nums(i) + nums(j)) == target) =>
            (i, j)
        }
    } collectFirst {
      case Some(tuple) => tuple
    }

Which you can use like:

twoSum(ArraySeq(1, 3, 5, 2, 0, 4), target = 8)
// res: Option[(Int, Int)] = Some((1,2))

twoSum(ArraySeq(1, 3, 5, 2, 0, 4), target = 10)
// res: Option[(Int, Int)] = None
Related