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

【问题标题】：Avoiding loops or list comprehension with numpy使用 numpy 避免循环或列表理解
【发布时间】：2020-10-14 15:04:03
【问题描述】：

是否可以更换

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 函数可以比这个列表理解提供性能提升？

【问题讨论】：

如果y=2 np.concatenate([np.where(x == i)[0] for i in range(y)]) 不会返回 [0, 3, 1, 2, 4, 5] 因为range(y) 给出了列表的可迭代[0, 1]

【解决方案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])

【讨论】：

argsort 的问题在于，如果值相同，它似乎不一定按照它们在数组中出现的顺序返回索引。例如。 np.argsort([0,1,0,0]) 不能保证返回 [0,2,3,1] - 它可以返回 [0,3,2,1]
实际上使用“稳定”算法进行排序（np.argsort 中的 kind="stable"）似乎以所需的方式保持顺序。感谢您的回答

【解决方案2】：

使用argsort 会起作用。

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

【讨论】：