CODEWITHSUNDEEP
pythonlistlist-comprehension
alexvassel
alexvassel

Reputation: 10740

Python, working with list comprehensions

I have such code:

a = [[1, 1], [2, 1], [3, 0]]

I want to get two lists, the first contains elements of 'a', where a[][1] = 1, and the second - elements where a[][1] = 0. So 

first_list = [[1, 1], [2, 1]] 

second_list = [[3, 0]].

I can do such thing with two list comprehension:

first_list = [i for i in a if i[1] == 1]

second_list = [i for i in a if i[1] == 0]

But maybe exists other (more pythonic, or shorter) way to do this? Thanks for your answers.

Upvotes: 5

Views: 393

Answers (6)

Rusty Rob
Rusty Rob

Reputation: 17173

list comprehensions are great. If you want slightly more simple code (but slightly longer) then just use a for loop.

Yet another option would be filters and maps:

a = [[1, 1], [2, 1], [3, 0]]
g1=filter(lambda i: i[1]==1,a)
g1=map(lambda i: i[0],g1)
g2=filter(lambda i: i[1]==0,a)
g2=map(lambda i: i[0],g2)
print g1
print g2

Upvotes: 0

Duncan
Duncan

Reputation: 95622

If the lists are reasonably short then two list comprehensions will do fine: you shouldn't be worried about performance until your code is all working and you know it is too slow.

If your lists are long or the code runs often and you have demonstrated that it is a bottleneck then all you have to do is switch from list comprehensions to a for loop:

first_list, second_list = [], []
for element in a:
    if element[1] == 1:
        first_list.append(element)
    else:
        second_list.append(element)

which is both clear and easily extended to more cases.

Upvotes: 0

sunqiang
sunqiang

Reputation: 6492

what about this, 

In [1]: a = [[1, 1], [2, 1], [3, 0]]

In [2]: first_list = []

In [3]: second_list = []

In [4]: [first_list.append(i) if i[1] == 1 else second_list.append(i) for i in a]
Out[4]: [None, None, None]

In [5]: first_list, second_list
Out[5]: ([[1, 1], [2, 1]], [[3, 0]])

instead of two sublist, I prefer dict (or defaultdict, OrderedDict, Counter, etc.)

In [6]: from collections import defaultdict

In [7]: d = defaultdict(list)

In [8]: [d[i[1]].append(i) for i in a]
Out[8]: [None, None, None]

In [9]: d
Out[9]: {0: [[3, 0]], 1: [[1, 1], [2, 1]]}

Upvotes: 0

FlaPer87
FlaPer87

Reputation: 1274

If you want to have it in a single line you could do something like

first_list, second_list = [i for i in a if i[1] == 1], [i for i in a if i[1] == 0]

Remember that, "Explicit is better than implicit."

Your code is fine

Upvotes: 3

Ignacio Vazquez-Abrams
Ignacio Vazquez-Abrams

Reputation: 798526

You can use sorted() and itertools.groupby() to do this, but I don't know that it would qualify as Pythonic per se:

>>> dict((k, list(v)) for (k, v) in itertools.groupby(sorted(a, key=operator.itemgetter(1)), operator.itemgetter(1)))
{0: [[3, 0]], 1: [[1, 1], [2, 1]]}

Upvotes: 2

Björn Pollex
Björn Pollex

Reputation: 76778

List comprehension are very pythonic and the recommended way of doing this. Your code is fine.

Upvotes: 7

Related Questions