具有“行”和索引的 ismember 的 Python 版本

已经提出了类似的问题，但没有一个答案完全满足我的需要 - 有些答案允许多维搜索（又名 matlab 中的“行”选项），但不返回索引。有些返回索引但不允许行。我的数组非常大（1M x 2），并且我已经成功地创建了一个有效的循环，但显然这非常慢。在matlab中，内置的ismember函数大约需要10秒。

这是我正在寻找的：

a=np.array([[4, 6],[2, 6],[5, 2]])

b=np.array([[1, 7],[1, 8],[2, 6],[2, 1],[2, 4],[4, 6],[4, 7],[5, 9],[5, 2],[5, 1]])

执行此操作的具体 matlab 函数是：

[~,index] = ismember(a,b,'rows')

where

index = [6, 3, 9] 

import numpy as np

def asvoid(arr):
    """
    View the array as dtype np.void (bytes)
    This views the last axis of ND-arrays as bytes so you can perform comparisons on
    the entire row.
    http://stackoverflow.com/a/16840350/190597 (Jaime, 2013-05)
    Warning: When using asvoid for comparison, note that float zeros may compare UNEQUALLY
    >>> asvoid([-0.]) == asvoid([0.])
    array([False], dtype=bool)
    """
    arr = np.ascontiguousarray(arr)
    return arr.view(np.dtype((np.void, arr.dtype.itemsize * arr.shape[-1])))


def in1d_index(a, b):
    voida, voidb = map(asvoid, (a, b))
    return np.where(np.in1d(voidb, voida))[0]    

a = np.array([[4, 6],[2, 6],[5, 2]])
b = np.array([[1, 7],[1, 8],[2, 6],[2, 1],[2, 4],[4, 6],[4, 7],[5, 9],[5, 2],[5, 1]])

print(in1d_index(a, b))

prints

[2 5 8]

这相当于 Matlab 的 [3,6,9]，因为 Python 使用基于 0 的索引。

一些注意事项：

1. 索引按升序返回。他们不对应 到项目的位置a in b.
2. asvoid 适用于整数数据类型，但使用 asvoid 时要小心 在浮点数据类型上，因为asvoid([-0.]) == asvoid([0.])回报array([False]).
3. asvoid 在连续数组上效果最好。如果数组不连续，数据将被复制到连续数组，这会降低性能。

尽管有警告，人们可能会选择使用in1d_index无论如何，为了速度：

def ismember_rows(a, b):
    # http://stackoverflow.com/a/22705773/190597 (ashg)
    return np.nonzero(np.all(b == a[:,np.newaxis], axis=2))[1]

In [41]: a2 = np.tile(a,(2000,1))
In [42]: b2 = np.tile(b,(2000,1))

In [46]: %timeit in1d_index(a2, b2)
100 loops, best of 3: 8.49 ms per loop

In [47]: %timeit ismember_rows(a2, b2)
1 loops, best of 3: 5.55 s per loop

So in1d_index速度快了约 650 倍（对于长度在数千以内的数组），但再次注意，这种比较并不完全是同类比较，因为in1d_index按升序返回索引，同时ismember_rows返回顺序行中的索引a出现在b.