1. Overview
In this tutorial, we’ll see how to obtain the first N elements of a Kotlin Array or Collection. We’ll first cover how to do this by iterating through an array and storing the values in a temporary variable. We’ll then see how to leverage the take() and takeWhile() functions, which are built into Kotlin. These will help us achieve the same results with less code.
2. The Classic, Iterative Way
One way of doing this is to implement this function manually. The classic way to do that is iterating through the original array and storing values in a temporary variable, then returning this temporary variable as the result, discarding the rest. Let’s write a unit test to demonstrate this concept:
@Test
fun `Given an array take first n elements`() {
val originalArray = intArrayOf(1, 2, 3, 4, 5)
val n = 3
val tempArray = IntArray(n)
for (i in 0 until n) {
tempArray[i] = originalArray[i]
}
assertContentEquals(intArrayOf(1, 2, 3), tempArray)
}
3. The Kotlin Way with take()
This same problem is covered by the Kotlin Standard Library(built-in), and we can use provided functions to do exactly what we want, to the letter! From the Kotlin Docs, we see that the take() function “returns a list containing first n elements” — this definition aligns perfectly with our problem statement.
Using it is as simple as calling it on our original array. Let’s add another unit test to check for the same result:
@Test
fun `Given an array take first n elements using kotlin function`() {
val originalArray = intArrayOf(1, 2, 3, 4, 5)
val n = 3
val tempArray: List<Int> = originalArray.take(n)
assertContentEquals(listOf(1, 2, 3), tempArray)
}
A small detail that can go unnoticed: Instead of an Array, we are now dealing directly with a List. The difference is subtle, but the types are different. To understand the difference between the two, our article Difference Between List and Array in Kotlin might be helpful.
3.1. Taking Elements From a List
We can notice from the previous segment that Kotlin already prefers using Lists instead of Arrays. It’s reasonable to expect that the same functions available for an Array should be available for a List as well. By changing a bit of our previous unit test, we can see that the syntax works exactly as expected:
@Test
fun `Given a list take first n elements using kotlin function`() {
val originalList = listOf(1, 2, 3, 4, 5)
val n = 3
val tempList = originalList.take(n)
assertContentEquals(listOf(1, 2, 3), tempList)
}
3.2. The takeWhile() Function
An advanced version of take() is takeWhile(). From the Kotlin Docs, we learn that it “returns a list containing first elements satisfying the given predicate”.
Let’s experiment with this function on our list, to understand how we can call it:
@Test
fun `Given a list take elements while predicate`() {
val list = listOf(1, 2, 3, 4, 5)
val takenList = list.takeWhile { it < 4 }
assertContentEquals(takenList, listOf(1, 2, 3))
}
We must be careful when using this function on unsorted collections, as takeWhile() will go iteratively from start to end until the condition is no longer satisfied. Let’s see what would happen with a different list:
@Test
fun `Given a list take elements while predicate in an unsorted list`() {
val list = listOf(5, 3, 2, 4, 1)
val takenList = list.takeWhile { it < 4 }
assertContentEquals(takenList, listOf())
}
We’ll take elements from the start of the list while it satisfies the predicate. As the start of the list won’t satisfy the predicate, we won’t take any elements from the list.
4. Converting Back to Array
If an Array is actually what we need, we can use the toTypedArray() function available on List:
val list: List<Int> = listOf(1, 2, 3, 4)
val array: Array<Int> = list.toTypedArray()
If we’re using lists of primitive types such as Int or Long, we might want to leverage more specific functions, such as toIntArray() or toLongArray().
For more information on this subject, check our article Convert a List into an Array in Kotlin.
5. Conclusion
In this tutorial, we learned how to write code to get the first N elements of an Array or a List. We first learned the classical implementation, by manually iterating through the collection. We then learned how to leverage the built-in take() and takeWhile() functions to get the same results, and lastly, we saw how to turn the result back into an Array. As always, the code for this tutorial is available over on GitHub.