使用numpy避免循环或列表理解

Question

是否可以替换

np.concatenate([np.where(x == i)[0] for i in range(y)])

有不涉及循环的东西？

我想获取一个数组x，例如[0、1、2、0、2、2]和数字y，例如在这种情况下为2，并输出数组[0、3、1、2、4、5]。 例如。对于数组中的每个整数，写出它们的索引位置，以使它们“按顺序”。

也许某种numpy函数可以通过这种列表理解来提高性能？

Answer 1

这是一种使用argsort的方法：

# settings
x = np.array([0, 1, 2, 0 , 2, 2])
y = 2

# sort the index
u = np.argsort(x)

# filter those that are larger than y
mask = x[u]<=y
u[mask]

输出：

array([0, 3, 1, 2, 4, 5])

Answer 2

使用argsort即可。

numpy.argsort([0, 1, 2, 0 , 2, 2])
=> array([0, 3, 1, 2, 4, 5])