CODEWITHSUNDEEP
kotlincollectionspartitioning
pixel
pixel

Reputation: 26441

Split a list into three in Kotlin

I have the following data class:

data class Foo(val a: Int = 0, val b: Int = 0)

I have a list of Foo's with the following structure:

[ Foo(a = 1), Foo(a = 2), ..., Foo(b = 22), Foo(a = 5), Foo(a = 6), ... ]

(a group of items with a's, then one b, then a's again)

I would like to split above list into three sub-lists like that:

  1. [ Foo(a = 1), Foo(a = 2), ...]
  2. [ Foo(b = 22) ]
  3. [ Foo(a = 5), Foo(a = 6), ...]
  1. sublist of elements that have non-zero a property
  2. list of one element that has non-zero b
  3. remaining sublist of elements that have non-zero a property

Is it possible to achieve using groupBy or partition?

Upvotes: 2

Views: 3303

Answers (2)

cactustictacs
cactustictacs

Reputation: 19514

So you want to 1) ignore the first Foos where a=0, 2) start collecting them when you see Foos where a is non-zero, 3) when you hit a Foo where a=0, put that in another list, because b will be non-zero, 4) start collecting non-zero a's again in a third list?

If that's what you want (this is an extremely specific thing you want and you haven't been clear about it at all) you could do it this way:

data class Foo(val a: Int, val b: Int)

val stuff = listOf(Foo(0,1), Foo(1,2), Foo(3,0), Foo(0, 4), Foo(0, 5), Foo(6, 1), Foo(0,7))

fun main(args: Array<String>) {
    fun aIsZero(foo: Foo) = foo.a == 0

    // ignore initial zero a's if there are any
    with(stuff.dropWhile(::aIsZero)) {
        val bIndex = indexOfFirst(::aIsZero)
        val listOne = take(bIndex)
        val listTwo = listOf(elementAt(bIndex))
        val listThree = drop(bIndex+1).filterNot(::aIsZero)
    
        listOf(listOne, listTwo, listThree).forEach(::println)
    }
}

You can't use partition or groupBy because your predicate depends on the value of a, but also on whether it happens to represent that one element you want to put in the b list, and for the others whether they appear before or after that b element. Which you don't know before you start processing the list.

You could mess around with indices and stuff, but honestly your use case seems so specific that it's probably better to just do it imperatively instead of trying to cram it into a functional approach.

Upvotes: 0

Utku &#214;zdemir
Utku &#214;zdemir

Reputation: 7725

It is not possible to do it via groupBy or partition, because it is not possible to check the past state in those operations. However, you can do it via a fold operation and using mutable lists. Not sure if it fits to your needs but here it goes:

    val input = listOf(Foo(a = 1), Foo(a = 2), Foo(b = 22), Foo(a = 5), Foo(a = 6))
    val output: List<List<Foo>> = input.fold(mutableListOf<MutableList<Foo>>(mutableListOf())) { acc, foo ->
        val lastList = acc.last()
        val appendToTheLastList =
            lastList.isEmpty() ||
                    (foo.a != 0 && lastList.last().a != 0) ||
                    (foo.b != 0 && lastList.last().b != 0)
        when {
            appendToTheLastList -> lastList.add(foo)
            else -> acc.add(mutableListOf(foo))
        }
        return@fold acc
    }

    println(output)

outputs:

[[Foo(a=1, b=0), Foo(a=2, b=0)], [Foo(a=0, b=22)], [Foo(a=5, b=0), Foo(a=6, b=0)]]

Note: I have to point out that this solution is not better than a solution with regular loops.

Upvotes: 1

Related Questions