[NumPy] searchsorted

np.searchsorted

sort가 된 기존의 ndarray A 에 대해,

입력으로 주어지는 ndarray B 의 element들의 값을 보고

기존의 ndarray A 의 어느 index에 놓이게 될지를 반환한다.

• 반환되는 ndarray의 shape는 query에 해당하는 B와 같다.
• 기본으로는 left를 기준으로 index를 반환한다. 같은 값이 이미 존재하는 경우, 해당 값의 위치의 left에 삽입하는 것으로 생각하고 index 반환.

import numpy as np

list_a = [1, 2, 3, 4, 5]
list_b = [2.5, 4]

print(np.searchsorted(list_a, list_b))

결과는 다음과 같다.

[2 3]

• A는 [1, 2, 3, 4, 5]이고, B는 [2.5, 4]이다.
• index는 0부터 시작하니 2.5는 A에서 index=2위치에 삽입될 수 있다. (sort상태를 유지하고 삽입한다.)
• 4의 경우, 동일한 값이 index=3 에 위치한다. left가 기본이므로 3이 반환된다.

---

다음과 같이 right를 인자로 줄 경우,

print(np.searchsorted(list_a, list_b, side='right'))

결과는 다음과 같다.

[2 4]

---

sort된 ndarray A를 가정하기 때문에, 정렬이 안된 A를 사용할 경우엔 제대로 된 반환값을 보장하지 않는다.
만일 정렬을 하지 않고 입력하려면, parameter sorter에 A 를 sorting하기 위한 index list를 넣어주면 된다.

import numpy as np

list_a = [8, 2, 3, 4, 5]
list_b = [2.5, 4]

print(np.searchsorted(list_a, list_b, sorter=np.argsort(list_a)))

결과는 다음과 같다.

[1 2]

pytorch에서도 torch.searchsorted로 제공된다.

---

같이 보면 좋은 자료들

https://ds31x.tistory.com/344

[Summary] NumPy(Numerical Python)