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

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?

CodePudding user response：

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]

CodePudding user response：

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)

CodePudding user response：

Perhaps the most elegant solution is

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

but probably not the fastest