mariogarcc
mariogarcc

Reputation: 388

Unzipping values inside list comprehension

I want to do the following:

A, B = [[x, y] for (x, y) in Z]

desiring an output:

A = [x1, x2, ...] # xN are the N x values in Z
B = [y1, y2, ...]

There is the obvious solution to do:

A, B = [x for (x, y) in Z], [y for (x, y) in Z]

but the actual code has a pretty big list comprehension with some conditions inside, so it should be kept that way to not only save lines of code (and performance!) but also achieve the desired program behaviour (actual current code only retrieves x for (x, y) in Z and assigns it to A).

Upvotes: 2

Views: 148

Answers (3)

Olivier Melançon
Olivier Melançon

Reputation: 22314

If you say there are multiple, possibly complex, conditions, then remember that using a for-loop may be more adapted.

A = []
B = []
for x, y in Z:
    if conditon_on_x:
        A.append(x)
    if condition_on_y:
        B.append(y)

Upvotes: 1

Silas Coker
Silas Coker

Reputation: 500

You can unzip by using zip, by expanding the variable like zip(*var):

>>> list_of_tuples = [('a', 0), ('b', 1), ('c', 2), ('d', 3), ('e', 4)]
>>> alpha, num = zip(*list_of_tuples)
>>> print(alpha)
('a', 'b', 'c', 'd', 'e')
>>> print(num)
(0, 1, 2, 3, 4)

This works because zip(*list_of_tuples) is expanded to zip( ('a',0), ('b',1), ('c',2), ('d',3), ('e',4) ), and zip zips all those tuples together, resulting, ironically enough, in an unzip.

But I don't think doing this with a single list comprehension is actually possible, sorry!

Upvotes: 1

user8105524
user8105524

Reputation:

A, B = zip(*([x, y] for (x, y) in Z))

Should work. Depending on the type of Z, you could probably get away with:

A, B = zip(*Z)

Upvotes: 3

Related Questions