用于列表的python numpy.where函数

时间:2013-08-12 13:46:31

标签: python list numpy where

我有一个包含两个列表的python程序,如下例所示:

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

是否有一种优雅的方法来创建第三个列表,其中包含list2在list2为1的位置的元素?我正在寻找类似于数组的numpy.where函数或更好的方式:

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

是否可以创建一个与array3等效的list3,在此示例中包含“b”和“c”,或者我是否必须强制转换或使用循环?

3 个答案:

答案 0 :(得分:2)

这正是itertools.compress所做的。

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

答案 1 :(得分:1)

你可以在这里使用列表理解。

>>> 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']

或者使用,

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

答案 2 :(得分:0)

您可以执行以下列表理解:

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

您的示例给出了:

['b','c']

构建get_val_where函数:

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

使用get_val_where(list1,list2,1)

之类的内容