Is there a nice way to chunk in Kotlin? I need a list into chunks of consecutive elements whose sum is less than X.

Example [1, 2, 3, 4, 5, 6, 7, 8, 9, 10 15 13 17 10] -> [[1, 2, 3, 4, 5], [6, 7], [8], [9], [10]] (each chunk has sum less than 16)

I thought maybe chunked method would work but I can chunk only by size. For now, I just made the for loop in an imperative way.

CodePudding user response：

After some testing, this problem has two main aspects.

• Create a new collection with consecutive numbers only
• Map consecutive numbers into list of lists

For the former part, a while-loop is appropriate because it is easy to exit when condition is no longer met. That is not straight forward with e.g. filter.

For the latter part, fold may be used to map (reduce) the list of consecutive numbers into a list of lists.

Code should handle cases where a new sequence of consecutive numbers appear in the input, e.g.:

val input = listOf(1, 2, 3, 4, 5, 8, 6, 7)
val expected = listOf(listOf(1, 2, 3, 4, 5))

And when numbers that are greater than 15 appears in input

val input = listOf(15, 16)
val expected = listOf(listOf(15))

Code is written inside a Kotest.

import io.kotest.core.spec.style.BehaviorSpec
import io.kotest.matchers.shouldBe

class MyTest : BehaviorSpec({

    given("some input") {
        val input = listOf(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 13, 17, 10)

        val expected = listOf(
            listOf(1, 2, 3, 4, 5),
            listOf(6, 7),
            listOf(8),
            listOf(9),
            listOf(10)
        )

        `when`("folding") {

            val result = mutableListOf(input[0]).apply {
                /** build a list of consecutive numbers only */
                var index = 1
                while (index < input.size
                    && input[index] < 16
                    && last() == input[index] - 1
                ) {
                    add(input[index  ])
                }
            }.fold(mutableListOf(mutableListOf<Int>())) { accumulator, current ->
                accumulator.apply {
                    /** if sum of current "chunk" will be greater than 15 */
                    if (last().sum()   current > 15) {
                        /** add a new "chunk" with current element */
                        add(mutableListOf(current))
                    } else {
                        /** add element to current "chunk" */
                        last().add(current)
                    }
                }
            }

            then("result should be as expected") {
                result shouldBe expected
            }
        }
    }
})

apply: Extension method that uses this internally and returns the receiver, as opposed to run that returns the result of the lambda.

also and let are the counterparts to apply and run, but they are using it internally.

CodePudding user response：

Here's a general-purpose chunkedBy() function I wrote to achieve this:

/**
 * Splits a collection into sublists not exceeding the given size.
 * This is a generalisation of [List.chunked]; but where that 
 * limits the _number_ of items in each sublist, this limits 
 * their total size, according to a given sizing function.
 *
 * @param maxSize None of the returned lists will have a total size 
 *                greater than this (unless a single item does).
 * @param includeAll Whether to include all items, even those
 *                   that are individually over the limit.
 * @param size Function giving the size of an item.
 */
fun <T> Iterable<T>.chunkedBy(maxSize: Int, includeAll: Boolean = true, size: (T) -> Int): List<List<T>> {
    val result = mutableListOf<List<T>>()
    var sublist = mutableListOf<T>()
    var sublistSize = 0L
    for (item in this) {
        val itemSize = size(item)
        if (itemSize > maxSize && !includeAll)
            continue
        if (sublistSize   itemSize > maxSize && sublist.isNotEmpty()) {
            result  = sublist
            sublist = mutableListOf()
            sublistSize = 0
        }
        sublist.add(item)
        sublistSize  = itemSize
    }
    if (sublist.isNotEmpty())
        result  = sublist

    return result
}

// Example:
val list = listOf(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 13, 17, 10)
println(list.chunkedBy(15, false){ it })
// prints: [[1, 2, 3, 4, 5], [6, 7], [8], [9], [10], [15], [13], [10]]

(includeAll defaults to true, as that suited all the cases in which I've needed it. But setting it to false does what you want here.)

(I've used this in a few places, and it's come up on this site before — it seems useful enough that I'm surprised it's not in the standard library.)