Numpy数组条件索引查找：求类似Matlab find函数的实现方法

嘿，我太懂你习惯了Matlab的find函数，转用NumPy时找不到对应方法的别扭感！其实NumPy里不需要专门的find函数，用内置的方法就能轻松实现你要的功能，甚至更简洁。

针对你给出的示例需求，我给你两种常用的实现方式：

方法1：使用np.where()（最接近Matlab find的用法） #

np.where()可以返回满足条件的元素索引，对于一维数组，它返回的是一个包含索引数组的元组，我们取第一个元素即可拿到索引：

import numpy as np
x = np.linspace(-5,5,11)
# 获取所有x<0的元素索引
k_neg = np.where(x < 0)[0]
# 获取所有x>=0的元素索引
k_pos = np.where(x >= 0)[0]
y = np.zeros(len(x))
y[k_neg] = -1
y[k_pos] = 1

方法2：直接使用布尔索引赋值（更高效简洁） #

NumPy数组支持直接用布尔条件作为索引，不需要先提取索引数组，这种方式代码更短，执行效率也更高：

import numpy as np
x = np.linspace(-5,5,11)
y = np.zeros(len(x))
# 直接给满足x<0的位置赋值-1
y[x < 0] = -1
# 直接给满足x>=0的位置赋值1
y[x >= 0] = 1

如果是多维数组的场景，np.where()同样适用，比如二维数组arr，np.where(arr > 0)会返回两个数组，分别对应满足条件元素的行索引和列索引，用法和一维类似。

内容的提问来源于stack exchange，提问作者John