13

Imagine that I have an order list of tuples:

s = [(0,-1), (1,0), (2,-1), (3,0), (4,0), (5,-1), (6,0), (7,-1)]

Given a parameter X, I want to select all the tuples that have a first element equal or greater than X up to but not including the first tuple that has -1 as the second element.

For example, if X = 3, I want to select the list [(3,0), (4,0)]

One idea I had is: Get the cut-off key with

E = min (x [0] for x in s if (x [0] >= X) and (x [1] == -1) )

Then select elements with keys between the X and E:

R = [x for x in s if X <= x [0] < E]

That gives me what I want in R, but it seems really inefficient, involving two table scans. I could do it in a for loop, discarding tuples with keys too small, and break when I hit the first blocking tuple. But for runs like a dog compared to list selection.

Is there a super-efficient, python-esque (2.7) way of doing this?

martineau
  • 119,623
  • 25
  • 170
  • 301
Terrible Tadpole
  • 607
  • 6
  • 20

2 Answers2

28

You can simply filter the tuples from the list as a generator expression and then you can stop taking the values from the generator expression when you get the first tuple whose second element is -1, like this

>>> s = [(0,-1), (1,0), (2,-1), (3,0), (4,0), (5,-1), (6,0), (7,-1)]
>>> from itertools import takewhile
>>> X = 3
>>> list(takewhile(lambda x: x[1] != -1, (item for item in s if item[0] >= X)))
[(3, 0), (4, 0)]

Here, the generator expression, (item for item in s if item[0] >= X) will give values one-by-one, on demand, (they are not generated all at once, so we save memory here) which are greater than or equal to X.

Then, we take values from that generator expression, only till we find a tuple whose second element is not equal to -1, with itertools.takewhile.

thefourtheye
  • 233,700
  • 52
  • 457
  • 497
1

Here's a slightly hacky to implement takewhile as part of a generator expression:

def stop(): raise StopIteration

result = (
    stop() if item[1] == -1 else
    item
    for item in s
    if item[0] >= X
)

Or with different phrasing:

def until(cond):
    if cond: raise StopIteration
    return True

result = (
    item for item in s
    if item[0] >= X
    and until(item[1] == -1)
)
Eric
  • 95,302
  • 53
  • 242
  • 374
  • Interestingly enough, similar approach won't work in list comprehensions, because the `if` clause is outside the `StopIteration` handling by the comprehension/generator internal code, so the exception propagates. (With generators it does too, but it gets handled naturally by the code that *uses* the generator). This comment is for potential benefit of those who would try to apply this technique to a comprehension and get surprized, like I was. :) – atzz Mar 22 '16 at 22:31