output nested list with list comprehension

Question

I have the following list:

 a = [[0,11], [1,12], [3,14], [5,16],[7,18]]

I want to create a new list b that takes all elements satistying

 a[:,0] > 3

and prepend the index in a of the selected nested list. So b shoudl look like:

b = [[3, 5, 16],[4, 7,18]]

With:

b = [points, points in a if points[0]>3]

will output b = [[5, 16],[7,18]] with the index, and

b = [[index, points], index, points in enumerate(a) if points[0]>3]

shows an error. How can I generated b using list comprehension?

Answer 1

You can use a tuple of throwaway variables for holding the nested lists items:

>>> [[i,j,k] for i,(j,k) in enumerate(a) if j>3]
[[3, 5, 16], [4, 7, 18]]

Or as a more comprehensive approach for lists with more items you can use following list comprehension:

>>> [[i] + points for i, points in enumerate(a) if points[0]>3]
[[3, 5, 16], [4, 7, 18]]

Python 3 version using unpacking assignment:

>>> [[i, first] + rest for i,(first, *rest) in enumerate(a) if first>3]
[[3, 5, 16], [4, 7, 18]]

Answer 2

If a is a NumPy array (which it seems to be in your actual use case), you could do the following:

# One of the rare cases when numpy.where isn't superfluous.
indices = numpy.where(a[:, 0] > 3)[0]

index_column = indices[:, None]
selected_rows = a[indices]

b = numpy.hstack([index_column, selected_rows])

Or with less intermediate variables:

indices = numpy.where(a[:, 0] > 3)[0]
b = numpy.hstack([indices[:, None], a[indices]])

For large sizes of a , this will probably outperform solutions based on enumerate or other Python-level iteration techniques.

Answer 3

好吧，我开始回答这个问题，但是其他人早就发布了，所以这是另一种方法（没有列表理解，但仍然很神秘）：

map(lambda y: [y[0]] + y[1], filter(lambda x: x[1][0] > 3, enumerate(a)))