CODEWITHSUNDEEP
pythonlist
Gabriel
Gabriel

Reputation: 42329

Create new list based on values taken from sublist

I have a list composed of a given number of sublists. I need to create a new (smaller) list by rejecting items in all sublists based on the values stored in the first one. This is what a MWE of what I mean:

a = [[0,0,0,1,1,1,2,2,2], [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9],
[0.21,0.22,0.23,0.24,0.25,0.26,0.27,0.28,0.29]]

b = [[] for _ in range(len(a))]
for indx, item in enumerate(a[0]):
    if item == 2:
        b[0].append(a[0][indx])
        b[1].append(a[1][indx])
        b[2].append(a[2][indx])

which results in:

[[2, 2, 2], [0.7, 0.8, 0.9], [0.27, 0.28, 0.29]]

This works, but I'm looking for a more pythonic way of creating my b list.

Upvotes: 1

Views: 1307

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121534

Use zip() to combine the lists into columns, filter the columns, then use zip() again to transpose back to rows:

b = zip(*(col for col in zip(*a) if col[0] == 2))

Demo:

>>> a = [[0,0,0,1,1,1,2,2,2], [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9],
... [0.21,0.22,0.23,0.24,0.25,0.26,0.27,0.28,0.29]]
>>> zip(*(col for col in zip(*a) if col[0] == 2))
[(2, 2, 2), (0.7, 0.8, 0.9), (0.27, 0.28, 0.29)]

This creates a list of tuples. Should you absolutely need lists, map to lists:

b = map(list, zip(*(col for col in zip(*a) if col[0] == 2)))

or, in Python 3, use a list comprehension:

b = [list(t) for t in zip(*(col for col in zip(*a) if col[0] == 2))]

Upvotes: 4

Related Questions