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

Question

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?

An automated way to obtain the indices of the values of w into v. A variable idx_new such that v[idx_new] == w is True element by element — G. Gare
– G. Gare, Commented Jul 28, 2022 at 12:49
@G.Gare Doing searchsorted in the array after inserting the elements is worst than sorting the elements beforehand — Dani Mesejo
– Dani Mesejo, Commented Jul 28, 2022 at 13:09

Dani Mesejo · Accepted Answer · 2022-07-28 12:53:57Z

2

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]

edited Jul 28, 2022 at 12:53

answered Jul 28, 2022 at 12:52

Dani Mesejo

62.2k6 gold badges56 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

10 Comments

G. Gare Over a year ago

doesn't work if w is not sorted

Dani Mesejo Over a year ago

@G.Gare Why don't sort w then?

G. Gare Over a year ago

if w has a million entries you may want to avoid that, but it's a good solution nevertheless!

Mad Physicist Over a year ago

You need to apply the argsort of the argsort of the index to the range for a general solution

G. Gare Over a year ago

This answer is the best if the vector w is sorted, or with a pre-sorting of w

|

Giovanni Tardini · Accepted Answer · 2022-07-28 13:52:05Z

2

Perhaps the most elegant solution is

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

but probably not the fastest

answered Jul 28, 2022 at 13:52

Giovanni Tardini

5683 silver badges7 bronze badges

1 Comment

Mad Physicist Over a year ago

This is the second to most elegant solution :) Most elegant is pre-sorting the data to begin with, since that only requires one sort instead of 2

Giovanni Tardini · Accepted Answer · 2022-07-28 13:05:05Z

1

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)

answered Jul 28, 2022 at 13:05

Giovanni Tardini

5683 silver badges7 bronze badges

Collectives™ on Stack Overflow

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

3 Answers 3

10 Comments

1 Comment

Comments

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

10 Comments

1 Comment

Comments

Related