在具有非数字项的非唯一数组中查找唯一数组的索引

Find indices of unique array in a non-unique array with non-numeric items

我有以下数组:

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

实际存储在数组中的数据是uuid,例如:

123e4567-e89b-12d3-a456-426614174000

我想在 a 中搜索 b 并获取索引:

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

这个解决方案对我有用:

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

但问题是我正在处理巨大的数组(非唯一子数组中有 15M,str_ 类型的唯一子数组中有 150k),因此对于给定的操作,我需要 1.8TB 的内存没有。

知道如何解决此问题或使用我自己的解决方案解决内存限制吗?

谢谢。

由于索引的顺序并不真正相关,您可以在结果上使用 np.isin, and then np.flatnonzero 来检索返回数组为 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)*len(b))) 不同,这应该相当快且内存效率高 (O(len(a))),即使使用问题中提到的数组大小也是如此:

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)

这是一个基于 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)