Numpy: Эффективные элементы массива A с точки зрения индексов совпадающих элементов массива B

#python #arrays #numpy

Вопрос:

Предположим, у нас есть некоторый массив numpy A = [1, 2, 2, 3] и его уникальный массив B = [1, 2, 3] , где массив B в основном является набором для массива A.

Я хочу получить новый массив C = [0, 1, 1, 2] , который является массивом A, но его элементы заменены индексами соответствующих элементов в массиве B.

Как бы я этого добился?

Примечание: Мои массивы A и B содержат значения 18m и 138k соответственно, поэтому решение должно быть эффективным.

Ответ №1:

С помощью numpy.where :

 A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

C = np.where(A==B[:,None])[0]
 

выход: array([0, 1, 1, 2])

Ответ №2:

Воспользуйся:

 import numpy as np

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

res = (A[:, None] == B).argmax(1)
print(res)
 

Выход

 [0 1 1 2]
 

Альтернатива:

 _, res = (A[:, None] == B).nonzero()
print(res)
 

Выход

 [0 1 1 2]
 

Альтернатива, которая должна использовать меньше памяти:

 import numpy as np

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

# find the uniques values together with the inverse
Bp, inverse = np.unique(A, return_inverse=True)

# find the indices mapping Bp to B
indices = (Bp[:, None] == B).argmax(1)

# index on the mapping
result = indices[inverse]
print(result)
 

Выход

 [0 1 1 2]
 

Если B всегда есть уникальные отсортированные элементы A, вы можете сделать это напрямую:

 import numpy as np

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

# find the uniques values together with the inverse
_, inverse = np.unique(A, return_inverse=True)

print(inverse)
 

Выход

 [0 1 1 2]
 

Комментарии:

1. NB. np.where(condition) на самом деле это короткий путь для np.asarray(condition).nonzero() 😉

2. Спасибо! Решение сработало, но мне интересно, есть ли более эффективный способ выполнить эту операцию? Длина моих A и B составляет около 18 м и 138 тыс. Каждый; Я делаю это в Google Colab, и он сжигает оперативную память.

3. @dangtony98 dos B имеет уникальные элементы?

4. ДА. B обладает всеми уникальными элементами A.

5. @dangtony98 Они отсортированы? Как вы получаете эти два массива?