CODEWITHSUNDEEP
pythonnumpy
G. Gare
G. Gare

Reputation: 267

NumPy: Get indices of elements of array after insertion in sorted array

Consider the code

import numpy as np

v = np.linspace(0, 9, 10)
w = np.array([3.5, 4.5])
idx = np.searchsorted(v, w)
v = np.insert(v, idx, w)

print(idx, v[idx])

which outputs

[4 5] array([3.5, 4. ])

The variable idx contains the indices of the elements of w if they were inserted in v one by one. When inserting an array into another like above, only the value of idx corresponding to the minimum value of w will give as well the position of the same value into v.

Is there a way with numpy functions to obtain the indices of the elements of w once inserted?

Upvotes: 0

Views: 144

Answers (3)

Giovanni Tardini
Giovanni Tardini

Reputation: 568

Perhaps the most elegant solution is

idx_new = idx + np.argsort(np.argsort(idx))

but probably not the fastest

Upvotes: 2

Giovanni Tardini
Giovanni Tardini

Reputation: 568

I tried this, I hope it is general enough

import numpy as np
v = np.linspace(0, 9, 10)
w = np.array([3.5, 4.5])
idx = np.searchsorted(v, w)
v = np.insert(v, idx, w)

print(idx, v[idx])
idx_new = np.searchsorted(v, w)
print(idx_new)
print(v[idx_new], w)

Upvotes: 1

Dani Mesejo
Dani Mesejo

Reputation: 61910

One solution:

import numpy as np

v = np.linspace(0, 9, 10)
w = np.array([3.5, 4.5])
idx = np.searchsorted(v, w)
v = np.insert(v, idx, w)

new_idx = idx + np.arange(len(idx))
print(new_idx, v[new_idx])

Output

[4 5] [3.5 4.5]

Upvotes: 2

Related Questions