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

Question

已经问过类似的问题，但没有一个答案完全符合我的需要 - 有些允许多维搜索（又名 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')

在哪里

index = [6, 3, 9] 

Answer 1

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))

印刷

[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

因此in1d_index 速度快了约 650 倍（对于长度在数千以下的数组），但再次注意，比较并不完全是同类比较，因为in1d_index按升序返回索引，而ismember_rows按 中a显示的行的顺序返回索引b。