Home > Software engineering >  Split a lists into multiple sub lists with predicate - looking for a nice way to write this code
Split a lists into multiple sub lists with predicate - looking for a nice way to write this code

Time:12-09

Lets say I have a list of integers: [2, 4, 4, 6, 1, 8, 3, 5, 2, 2, 2], and I want to split this lists into sub lists so each sub list contain values that match the predicate, or if it doesn't then only the value that didn't match, and it should be in the order of the original list.

So for example, if the predicate is "is even" I expect to get the following result: [[2, 4, 4, 6],[1],[8],[3, 5],[2, 2, 2]]

I have a code that works, but it's pretty ugly in my opinion and I feel like there should be a much nicer way to write it:

private fun <T : Any> toSubLists(values: List<T>, predicate: (T) -> Boolean): List<List<T>> {
  val subLists = mutableListOf<List<T>>()
  var currentList = mutableListOf<T>()
  values.forEach {
    if (!predicate(it)) {
      if (currentList.isNotEmpty()) {
        subLists.add(currentList)
        currentList = mutableListOf()
      }
      subLists.add(listOf(it))
    } else {
      currentList.add(it)
    }
  }
  if (currentList.isNotEmpty()) {
    subLists.add(currentList)
  }
  return subLists
}

CodePudding user response:

Not sure if it's better because it might be harder to understand, but this is a way to do it

private fun <T : Any> toSubLists(values: List<T>, predicate: (T) -> Boolean): List<List<T>> = 
    values.withIndex().groupBy({
        var index = it.index
        while (index != 0 && predicate(values[index]) && predicate(values[index - 1])) index--
        index
    }){
        it.value
    }.values.toList()

see it working here

CodePudding user response:

There is a feature request for a generic version of this that accepts other key selectors than predicates.

In the meantime, for your use case I think we can simplify a bit:

fun <T> List<T>.groupRunsBy(predicate: (T) -> Boolean): List<List<T>> {
    val result = mutableListOf<MutableList<T>>()
    var currentlyMatching: Boolean? = null
    forEach {
        val itemMatches = predicate(it)
        if (itemMatches == currentlyMatching) {
            result.last().add(it)
        } else {
            currentlyMatching = itemMatches
            result.add(mutableListOf(it))
        }
    }
    return result
}

fun main() {
    val list = listOf(2, 4, 4, 6, 1, 8, 3, 5, 2, 2, 2)
    
    println(list.groupRunsBy { it % 2 == 0 }) // [[2, 4, 4, 6], [1], [8], [3, 5], [2, 2, 2]]
}

Tested here: https://pl.kotl.in/OHVr7S6lX

This version can easily be changed to support any selector that returns non-null keys.

  • Related