CODEWITHSUNDEEP
pythongrouping
rdlu
rdlu

Reputation: 630

Python Equivalent to Ruby's chunk_while?

I have some interesting using cases in my Ruby code that I want to convert to Python. I think we can leverage in any lib, I mainly use pandas and numpy.

For instance, imagine that you have an array of timed events, these events have a timestamp and another properties (an object or a tuple).

I want a list/array of groups, where these groups are "contiguous" events, with a grace period of g units (in this case, time units).

In my Ruby code I use something like this:

grouped_events = events.chunk_while do |previous_event, next_event|
   next_event.timestamp <= previous_event.timestamp + grace_period
end

Since I'm not using only in timed events, but with anything that I can sort (so they're comparable somehow), I ask: there's a generic way, or a known lib that does this?

Upvotes: 4

Views: 247

Answers (1)

Aran-Fey
Aran-Fey

Reputation: 43246

Python doesn't have an equivalent function. You'll have to write your own.

Here's my implementation, using an iterator and the yield statement:

def chunk_while(iterable, predicate):
    itr = iter(iterable)

    try:
        prev_value = next(itr)
    except StopIteration:
        # if the iterable is empty, yield nothing
        return

    chunk = [prev_value]
    for value in itr:
        # if the predicate returns False, start a new chunk
        if not predicate(prev_value, value):
            yield chunk
            chunk = []

        chunk.append(value)
        prev_value = value

    # don't forget to yield the final chunk
    if chunk:
        yield chunk

Which can be used like so:

>>> list(chunk_while([1, 3, 2, 5, 5], lambda prev, next_: next_ <= prev + 2))
[[1, 3, 2], [5, 5]]

Upvotes: 6

Related Questions