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