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

Question

Is it possible to replace

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

with something that doesn't involve looping?

I want to take an array x, eg [0, 1, 2, 0 , 2, 2], and a number y, eg 2 in this case, and output an array [0, 3, 1, 2, 4, 5]. Eg for each integer in the array, write their index locations such that they're "in order".

Perhaps some kind of numpy function that can provide a performance boost over this list comprehension?

Answer 1

Here's an approach that uses 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]

Output:

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

Answer 2

Using argsort would work.

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