Индексирование чисел из 2 массивов

Рассмотрим два массива numpy

a = np.array(['john', 'bill', 'greg', 'bill', 'bill', 'greg', 'bill'])
b = np.array(['john', 'bill', 'greg'])

Как я могу создать третий массив

c = np.array([0,1,2,1,1,2,1])

Такая же длина, как a, представляющая индекс каждой записи a в массиве b?

Я вижу способ, перебирая элементы b как b[i] и проверяя np.where(a == b[i]), но задавался вопросом, может ли numpy выполнить это быстрее или лучше/меньше строк кода.

Ответ 1

Вот один из вариантов:

import numpy as np

a = np.array(['john', 'bill', 'greg', 'bill', 'bill', 'greg', 'bill'])
b = np.array(['john', 'bill', 'greg'])

my_dict = dict(zip(b, range(len(b))))

result = np.vectorize(my_dict.get)(a)

Результат:

>>> result
array([0, 1, 2, 1, 1, 2, 1])

Ответ 2

Сортировка - хороший вариант для векторизации с numpy:

>>> s = np.argsort(b)
>>> s[np.searchsorted(b, a, sorter=s)]
array([0, 1, 2, 1, 1, 2, 1], dtype=int64)

Если ваш массив a имеет m элементы и b имеет n, сортировка будет O (n log n) и поиск O (m log n), что неплохо, Решения на основе словаря должны амортизироваться линейно, но если массивы не огромны, цикл Python может сделать их медленнее, чем это. И решения на основе широковещательной передачи имеют квадратичную сложность, они будут работать только быстрее для очень малых массивов.


Некоторые тайминги с вашим образцом:

In [3]: %%timeit
   ...: s = np.argsort(b)
   ...: np.take(s, np.searchsorted(b, a, sorter=s))
   ...: 
100000 loops, best of 3: 4.16 µs per loop

In [5]: %%timeit
   ...: my_dict = dict(zip(b, range(len(b))))
   ...: np.vectorize(my_dict.get)(a)
   ...: 
10000 loops, best of 3: 29.9 µs per loop

In [7]: %timeit (np.arange(b.size)*(a==b[:,newaxis]).T).sum(axis=-1)
100000 loops, best of 3: 18.5 µs per loop

Ответ 3

Создайте словарь для перевода каждой строки в число, а затем используйте numpy.vectorize для создания выходного массива

>>> import numpy as np
>>> a = np.array(['john', 'bill', 'greg', 'bill', 'bill', 'greg', 'bill'])
>>> b = np.array(['john', 'bill', 'greg'])
>>> d = {k:v for v, k in enumerate(b)}
>>> c = np.vectorize(d.get)(a)
>>> c
 array([0, 1, 2, 1, 1, 2, 1])

Это более эффективно, чем цикл, и делает np.where(a == b[i]), потому что вы только один раз посещаете один элемент массива.

Ответ 4

Полноценное решение:

(arange(b.size)*(a==b[:,newaxis]).T).sum(axis=-1)

Ответ 5

Другое решение возможно:

arr, bSorted, ind =  np.unique(a, return_index=True, return_inverse=True)
c = bSorted[ind]

Если вы хотите получить уникальные элементы из a и не заботитесь о порядке в b, то есть b и, следовательно, c будет выглядеть по-другому, тогда его можно упростить до

b, c = np.unique(a, return_inverse=True)

Ответ 6

Поскольку массив b содержит уникальные элементы, равенство с элементом a может быть только когда-либо с одним элементом b. Если все элементы a определенно находятся в b, то

import numpy as np
indices = np.where(a[:, np.newaxis] == b)[1]

сделает трюк. Если вы не уверены, находятся ли все элементы a в b, то

in_b, indices = np.where(a[:, np.newaxis] == b)

будет собирать все элементы a, которые содержатся в b в in_b