Search code examples
kotlincollectionssequence

Mapping over a collection with a predicate including the first element on which predicate fails


I would like to map over a collection with a predicate such that:

  1. mapping should stop if the predicate is false
  2. the resulting collection must contain the element for which the predicate was false
  3. no additional mapping should occur after the predicate was false.

There is Sequence.takeWhile which satisfies 1 and 3 but not 2.

An example with takeWhile:

val seq = listOf(1, 2, 3, 4).asSequence()
seq.map { println("mapping: $it"); it }
    .takeWhile { it < 3 }
    .also { println(it.toList()) }    }

The output is

mapping: 1
mapping: 2
mapping: 3
[1, 2]

I'd need the result to be [1, 2, 3]


Solution

  • I don't think there is such a thing in the standard library. If you want to stick to standard library functions, you most likely need to iterate the sequence twice.

    This is something you can easily implement yourself:

    fun <T> Sequence<T>.inclusivelyTakeWhile(predicate: (T) -> Boolean) = sequence {
        for (t in this@inclusivelyTakeWhile) {
            yield(t)
            if (!predicate(t)) {
                break
            }
        }
    }
    

    Then doing:

    seq.map { println("mapping: $it"); it }
        . inclusivelyTakeWhile { it < 3 }
        .also { println(it.toList()) }  
    

    produces the desired result.

    Using the sequence function here probably adds some overhead, compared to the standard library's implementation of TakeWhileSequence. If you care about this, you can create a modified version of TakeWhileSequence and return that instead.