使用 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：# settingsx = np.array([0, 1, 2, 0 , 2, 2])y = 2# sort the indexu = np.argsort(x)# filter those that are larger than ymask = x[u]<=yu[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])