使用 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])