Athar
Athar

Reputation: 1063

Python: Can this be done in a single list comprehension statement

What I am trying to do is extracting zeroth element in a list and first element in another list of the given 2 dimensional list.

baseball = [[180, 78.4],
            [215, 102.7],
            [210, 98.5],
            [188, 75.2]]

x  =  [ a[0] for a in baseball ] 
y  =  [ a[1] for a in baseball ] 
print x
print y

Can this be done in a single list comprehension statement?

Upvotes: 2

Views: 81

Answers (2)

Padraic Cunningham
Padraic Cunningham

Reputation: 180401

If you don't mind tuples:

baseball = [[180, 78.4],
            [215, 102.7],
            [210, 98.5],
            [188, 75.2]]


x,y =  zip(*baseball)

If you really want lists:

x,y = map(list,zip(*baseball))

If you had more than two elements in each and wanted just certain elements like:

baseball = [[180, 1, 78.4],
            [215, 2, 102.7],
            [210, 3, 98.5],
            [188, 4, 75.2]]

from operator import itemgetter

x, y = zip(*map(itemgetter(0, 2), baseball))

That would give you:

 ((180, 215, 210, 188), (78.4, 102.7, 98.5, 75.2))

Upvotes: 2

Heman Gandhi
Heman Gandhi

Reputation: 1371

Assuming it's rectangular (ie. the length of the inner lists is consistent), you can implement the following:

def transpose(matrix):
    return [[matrix[j][i] for j in range(len(matrix))] for i in range(len(matrix[0]))]

Then, your problem is just a call to transpose (x, y = transpose(baseball)).

Upvotes: 2

Related Questions