python numpy.where function for lists

Question

I have a python program with two lists that look like the example below:

list1=["a","b","c","d"]
list2=[0,1,1,0]

Is there a elegant way to create a third list, that contains the elements of list1 at the position where list2 is 1? I am looking for something similar to the numpy.where function for arrays or better the elegant way:

array1=numpy.array(["a","b","c","d"])
array2=numpy.array([0,1,1,0])
array3=array1[array2==1]

Is it possible to create a list3 equivalent to array3, containing in this example "b" and "c" or do I have to cast or to use a loop?

Why not simply use array3.tolist()?

Daniel
– Daniel

2013-08-12 14:12:13 +00:00
Commented Aug 12, 2013 at 14:12 — Daniel
– Daniel, Commented Aug 12, 2013 at 14:12

Sukrit Kalra · Accepted Answer · 2013-08-12 14:50:12Z

2

You could use a list comprehension here.

>>> array1 = ["a", "b", "c", "d"]
>>> array2 = [0, 1, 1, 0]
>>> [array1[index] for index, val in enumerate(array2) if val == 1] # Or if val
['b', 'c']

Or use,

>>> [a for a, b in zip(array1, array2) if b]
['b', 'c']

edited Aug 12, 2013 at 14:50

answered Aug 12, 2013 at 13:48

Sukrit Kalra

34.7k7 gold badges73 silver badges71 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Abhijit · Accepted Answer · 2013-08-12 14:26:30Z

2

This is exactly what itertools.compress does.

>>> list1=["a","b","c","d"]
>>> list2=[0,1,1,0]
>>> import itertools
>>> list(itertools.compress(list1, list2))
['b', 'c']

answered Aug 12, 2013 at 14:26

Abhijit

64k20 gold badges143 silver badges209 bronze badges

Comments

Saullo G. P. Castro · Accepted Answer · 2013-08-12 14:03:36Z

0

You can do the following list comprehension:

[j for i,j in enumerate(list1) if list2[i]==1]

Which for your example gives:

['b','c']

Building a get_val_where function:

def get_val_where(list_val, list_where, val):
    return [j for i,j in enumerate(list_val) if list_where[i]==val]

To use like get_val_where(list1,list2,1)

edited Aug 12, 2013 at 14:03

answered Aug 12, 2013 at 13:58

Saullo G. P. Castro

59.4k28 gold badges191 silver badges244 bronze badges

Collectives™ on Stack Overflow

python numpy.where function for lists

3 Answers 3

Comments

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related