Рассмотрим два массива 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 выполнить это быстрее или лучше/меньше строк кода.
Вот один из вариантов:
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])
Сортировка - хороший вариант для векторизации с 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
Создайте словарь для перевода каждой строки в число, а затем используйте 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]), потому что вы только один раз посещаете один элемент массива.
Полноценное решение:
(arange(b.size)*(a==b[:,newaxis]).T).sum(axis=-1)
Другое решение возможно:
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)
Поскольку массив 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