Search code examples
scalascala-collections

Optimal way to find neighbors of element of collection in circular manner


I have a Vector and I'd like to find neighbors of given element.

Say if we have Vector(1, 2, 3, 4, 5) and then:

  • for element 2, result must be Some((1, 3))
  • for element 5, result must be Some((4, 1))
  • for element 1, result must be Some((5, 2))
  • for element 6, result must be None

and so on..

I have not found any solution in standard lib(please point me if there is one), so got the next one:

implicit class VectorOps[T](seq: Vector[T]) {
  def findNeighbors(elem: T): Option[(T, T)] = {
    val currentIdx = seq.indexOf(elem)
    val firstIdx = 0
    val lastIdx = seq.size - 1
    seq match {
      case _ if currentIdx == -1 || seq.size < 2 => None
      case _ if seq.size == 2 => seq.find(_ != elem).map(elem => (elem, elem))
      case _ if currentIdx == firstIdx => Some((seq(lastIdx), seq(currentIdx + 1)))
      case _ if currentIdx == lastIdx => Some((seq(currentIdx - 1), seq(firstIdx)))
      case _ => Some((seq(currentIdx - 1), seq(currentIdx + 1)))
    }
  }
}

The question is: how this can be simplified/optimized using stdlib?


Solution

  • def neighbours[T](v: Seq[T], x: T): Option[(T, T)] =
      (v.last +: v :+ v.head)
        .sliding(3, 1)
        .find(_(1) == x)
        .map(x => (x(0), x(2)))
    

    This uses sliding to create a 3 element window in the data and uses find to match the middle value of the 3. Adding the last/first to the input deals with the wrap-around case.

    This will fail if the Vector is too short so needs some error checking.


    This version is safe for all input

    def neighbours[T](v: Seq[T], x: T): Option[(T, T)] =
      (v.takeRight(1) ++ v ++ v.take(1))
        .sliding(3, 1)
        .find(_(1) == x)
        .map(x => (x(0), x(2)))