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

Question

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

mapping should stop if the predicate is false
the resulting collection must contain the element for which the predicate was false
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]

Sounds like something you can write your own extension function for. I don't think this exists in the standard library. — Sweeper, Apr 08 '22 at 11:57

score 1 · Accepted Answer · answered Apr 08 '22 at 12:17

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.

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

1 Answers1