Найдите индексы уникального массива в неуникальном массиве с нечисловыми c элементами

Question

У меня есть следующий массив:

import numpy as np
a = np.array(['a', 'b', 'c','a','a','d','e'])
b = np.array(['a','b'])

фактические данные, хранящиеся в массивах, - это uuids, например:

123e4567-e89b-12d3-a456-426614174000

Я хотел бы найти b в a и получить индексы :

array([0, 3, 4, 1])

это решение может работать для меня:

np.nonzero(b[:, None] == a)[1]

, но проблема в том, что я имею дело с огромными массивами (15M в неуникальных и 150k в уникальных подгруппах). -array with str_ type), поэтому для данной операции мне потребуется 1,8 ТБ памяти, которой у меня нет.

есть идеи, как я могу решить эту проблему или обойти ограничения памяти с помощью моего собственного решения?

спасибо.

Answer 1

Вот пример, основанный на view + lookup -

def map_indices_conststring(a, b):
    a2D = a.view(np.uint8)[::4].reshape(len(a),-1)
    b2D = b.view(np.uint8)[::4].reshape(len(b),-1)
    
    n = b2D.shape[1]
    lookup = np.zeros(256, dtype=bool)
    mask = np.ones(len(a), dtype=bool)
    for i in range(n):
        lookup[b2D[:,i]] = 1
        mask &= lookup[a2D[:,i]]
    out = np.flatnonzero(mask)
    return out

Пробный прогон -

In [46]: a
Out[46]: 
array(['a123', 'b232', 'c434', 'b235', 'a123', 'd223', 'b232'],
      dtype='<U4')

In [47]: b
Out[47]: array(['a123', 'b232'], dtype='<U4')

In [48]: map_indices_conststring(a, b)
Out[48]: array([0, 1, 4, 6])

Время для строковых данных с 1.5M неуникальным и 15K строковые массивы уникального размера -

In [2]: a = np.random.randint(10000000000,99999999999,(1500000)).astype(str)

In [3]: b = np.unique(np.random.randint(10000000000,99999999999,(15000)).astype(str))

In [4]: %timeit map_indices_conststring(a, b)
266 ms ± 2.63 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# @yatu's soln
In [5]: %timeit np.flatnonzero(np.isin(a,b))
1.03 s ± 3.75 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

Поскольку порядок индексов на самом деле не имеет значения, вы можете использовать np.isin, а затем np.flatnonzero в результатах, чтобы получить индексы, в которых возвращаемый массив is True:

a = np.array(['a', 'b', 'c','a','a','d','e'])
b = np.array(['a','b'])

np.flatnonzero(np.isin(a,b))
# array([0, 1, 3, 4], dtype=int64)

Это должно быть достаточно быстрым и эффективным с точки зрения памяти (O(len(a))), в отличие от широковещательного подхода (O(len(a)*len(b))), даже с размерами массива, упомянутыми в вопросе:

a = np.random.randint(0,15e2,int(15e6))
b = np.random.randint(0,150e3,int(150e3))

%timeit np.flatnonzero(np.isin(a,b))
# 2.58 s ± 28.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)