CODEWITHSUNDEEP
pythonperformancenumpy
brusselswanny
brusselswanny

Reputation: 40

Avoiding loops or list comprehension with numpy

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, e.g. [0, 1, 2, 0 , 2, 2], and a number y, e.g. 2 in this case, and output an array [0, 3, 1, 2, 4, 5]. E.g. 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?

Upvotes: 1

Views: 55

Answers (2)

shx2
shx2

Reputation: 64318

Using argsort would work.

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

Upvotes: 2

Quang Hoang
Quang Hoang

Reputation: 150735

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

Upvotes: 3

Related Questions