相当于 np.where 的“case”

Question

np.where 允许您选择要为布尔类型查询分配的值，例如

test = [0,1,2]
np.where(test==0,'True','False')
print test
['True','False','False']

这基本上是一个“if”语句。对于 numpy 数组，是否有一种Python式的方法来使用“if，else if，else”类型的语句（具有不同的情况）？

这是我的解决方法：

color = [0,1,2]
color = np.where(color==0,'red',color)
color = np.where(color==1,'blue',color)
color = np.where(color==2,'green',color)
print color
['red','blue','green']

但我想知道是否有更好的方法来做到这一点。

Answer 1

numpy.select()就是你想要的。这是 case when 的 numpy 版本。句法：

import numpy as np
color = np.array([0,1,2])
condlist = [color == 1, color == 2, color == 3]
choicelist = ['red', 'blue', 'green']
np.select(condlist, choicelist, default='unknown')

返回：

array(['unknown', 'red', 'blue'], dtype='<U7')

Answer 2

更 Pythonic 的方法之一是使用列表理解，如下所示：

>>> color = [0,1,2]
>>> ['red' if c == 0 else 'blue' if c == 1 else 'green' for c in color]
['red', 'blue', 'green']

如果您阅读它，就会非常直观。对于列表中的给定项目，如果颜色为，color则新列表中的值将为，否则为 。不过，我不知道列表理解中的 s 是否会超过 3。A'red'0'blue'1'green'if elsefor循环是合适的。

或者你可以使用字典，它可能更“Pythonic”，并且更具可扩展性：

>>> color_dict = {0: 'red', 1: 'blue', 2: 'green'}
>>> [color_dict[number] for number in color]
['red', 'blue', 'green']