Idiomatic scala way to implement this simple split-like function

Question

I'd like to implement a function that has the following characteristics:

• take an array (Array[A]) and a number of pieces to return (Int) as parameters
• returns the array array split in pieces pieces.
• all pieces should be of equal length if possible, else the first ones should be one element longer to the last ones (if the length of array isn't a multiple of pieces)

in Haskell I would have coded something along those lines:

split :: [a] -> Int -> [[a]]
split list pieces = go list (length list `div` pieces)
 where
  go xs n | l > n && m == 0 = take n xs : go (drop n xs) n
          | l > n           = take (n + 1) xs : go (drop (n + 1) xs) n
          | otherwise       = [xs]
   where
    l = length xs
    m = l `mod` n

Though in Scala I encounter many difficulties to code this (basic) function. For the recursion first, Array doesn't seem that adapted. Then, the if structure that'd allow me to implement the kind of guards I'm using in haskell aren't allowed in place of an expression, which seems weird to me. Another problem I have is that I wouldn't know how to make my scala code polymorphic (as my Haskell code is). Last but not least, I don't get how I'd go about making not only the a equivalent polymorphic, but also the Array instance, because there are many basic collections in Scala.

I'd be fine with either a complete explained solution or simple hints to answer my misunderstandings.

Answer 1

The translation is pretty straight forward. Instead of calling take and drop separately, I use splitAt which is potentially more efficient. Of course, you could make this tail-recursive. And I would also argue that "idiomatic" Scala might be using IndexedSeq.newBuilder inside the method, which is probably again more efficient, but I guess that's not what you are looking for.

Note that Scala's arrays are mutable, so you'd want Vector or the more general IndexedSeq instead.

import collection.immutable.{IndexedSeq => ISeq}

def split[A](seq: ISeq[A], pieces: Int): ISeq[ISeq[A]] = {
  val n = seq.size / pieces
  def loop(xs: ISeq[A]): ISeq[ISeq[A]] = {   
    val l = xs.size
    if(l > n) {
      val m = l % n
      val (begin, end) = xs.splitAt(if(m == 0) n else n + 1)
      begin +: loop(end)
    } else ISeq(xs)
  }
  loop(seq)
}

split(1 to 20, 3) // correct