Skip to main content
 首页 » 编程设计

python之一次获取 NumPy 数组中多行的索引

2025年05月04日14lidabo

我有 2 个数组,例如喜欢:

A: [[1 2 3][2 2 2][1 2 3][2 3 3][2 2 2][2 3 3][2 3 3]]   
B: [[1 2 3][2 2 2][2 3 3]] 

B 是 A 的已排序唯一行。
我需要:

C: [0 1 0 2 1 2 2] 

这是按 A 的顺序排列的 B 的索引列表。我想避免循环,因为即使对于非常大的数组,它也需要很快。

我找到的唯一解决方案仅适用于一维数组(例如 Getting the indices of several elements in a NumPy array at once)。
我认为这可以使用 np.void 以类似的方式解决:Find unique rows in numpy.array但我无法理解它:/

我需要在没有其他可用库的情况下使用 NumPy 1.10。

请您参考如下方法:

给定A和B,你可以生成C使用

In [25]: (B[:,None,:] == A).all(axis=-1).argmax(axis=0) 
Out[25]: array([0, 1, 0, 2, 1, 2, 2]) 

请注意,这假设 B 的每一行都在 A 中。 (否则,argmax 可能会返回等式为 False 的虚假索引。)


请注意,如果您使用的是 NumPy 1.13 或更新版本,那么 您可以使用 np.unique 同时生成 B 和 C:

In [33]: np.unique(A, axis=0, return_inverse=True) 
Out[33]:  
(array([[1, 2, 3], 
        [2, 2, 2], 
        [2, 3, 3]]), array([0, 1, 0, 2, 1, 2, 2])) 

请注意 Divakar's solution (使用 np.void)要快得多,尤其是当 A 有很多行时:

A = np.random.randint(10, size=(1000, 3)) 
B, C = np.unique(A, axis=0, return_inverse=True) 
 
In [44]: %%timeit 
   ....: A1D, B1D = view1D(A, B) 
   ....: sidx = B1D.argsort() 
   ....: out = argsort_unique(sidx)[np.searchsorted(B1D, A1D, sorter=sidx)] 
   ....:  
1000 loops, best of 3: 271 µs per loop 
 
In [45]: %timeit (B[:,None,:] == A).all(axis=-1).argmax(axis=0) 
100 loops, best of 3: 15.5 ms per loop