Python中的MATLAB的“ismember”函数的等价函数

Question

在多次尝试优化代码之后，最后一个资源似乎是尝试使用多个内核运行下面的代码。我不知道如何转换/重新构造我的代码，以便它可以使用多个核心运行得更快。如果我能得到实现最终目标的指导，我将不胜感激。最终目标是能够尽可能快地为数组A和B运行此代码，其中每个数组包含大约700000个元素。下面是使用小数组的代码。700k元素数组被注释掉。

import numpy as np

def ismember(a,b):
    for i in a:
        index = np.where(b==i)[0]
        if index.size == 0:
            yield 0
        else:
            yield index


def f(A, gen_obj):
    my_array = np.arange(len(A))
    for i in my_array:
        my_array[i] = gen_obj.next()
    return my_array


#A = np.arange(700000)
#B = np.arange(700000)
A = np.array([3,4,4,3,6])
B = np.array([2,5,2,6,3])

gen_obj = ismember(A,B)

f(A, gen_obj)

print 'done'
# if we print f(A, gen_obj) the output will be: [4 0 0 4 3]
# notice that the output array needs to be kept the same size as array A.

我要做的是模拟一个名为ismember[2]的MATLAB函数（格式为：[Lia,Locb] = ismember(A,B)）。我只想得到Locb部分。

From Matlab: Locb, contain the lowest index in B for each value in A that is a member of B. The output array, Locb, contains 0 wherever A is not a member of B

其中一个主要问题是，我需要能够尽可能有效地执行这个操作。为了测试，我有两个700k元素的数组。创建一个生成器并遍历生成器的值似乎不能快速完成任务。

Answer 1

在担心多核之前，我将使用字典消除ismember函数中的线性扫描：

def ismember(a, b):
    bind = {}
    for i, elt in enumerate(b):
        if elt not in bind:
            bind[elt] = i
    return [bind.get(itm, None) for itm in a]  # None can be replaced by any other "not in b" value

原始实现要求对a中的每个元素对B中的元素进行完全扫描，使其成为O(len(A)*len(B))。上述代码需要对B进行一次完整扫描才能生成dict Bset。通过使用dict，可以有效地查找a的每个元素的B常量中的每个元素，从而执行操作O(len(A)+len(B))。如果这仍然太慢，那么请担心使上述函数在多个内核上运行。

编辑：我也稍微修改了你的索引。Matlab使用0，因为它的所有数组都从索引1开始。Python/numpy从0开始数组，所以如果数据集是这样的

A = [2378, 2378, 2378, 2378]
B = [2378, 2379]

如果返回0表示没有元素，那么结果将排除A的所有元素。上面的例程返回None表示没有索引，而不是0。返回-1是一个选项，但Python会将其解释为数组中的最后一个元素。None如果用作数组的索引，则将引发异常。如果需要不同的行为，请将Bind.get(item,None)表达式中的第二个参数更改为要返回的值。

Answer 2

尝试使用列表理解

In [1]: import numpy as np

In [2]: A = np.array([3,4,4,3,6])

In [3]: B = np.array([2,5,2,6,3])

In [4]: [x for x in A if not x in B]
Out[4]: [4, 4]

通常，列表理解比循环理解快得多。

得到一个等长的列表

In [19]: map(lambda x: x if x not in B else False, A)
Out[19]: [False, 4, 4, False, False]

对于小数据集来说，这非常快：

In [20]: C = np.arange(10000)

In [21]: D = np.arange(15000, 25000)

In [22]: %timeit map(lambda x: x if x not in D else False, C)
1 loops, best of 3: 756 ms per loop

对于大型数据集，可以尝试使用multiprocessing.Pool.map()来加速操作。

Answer 3

斯夫斯图曼的出色回答很可能为你解决了这个问题。

我只想补充一点，你怎样才能在纽比实现同样的目标。

我使用了numpy的unique和in1d函数。

B_unique_sorted, B_idx = np.unique(B, return_index=True)
B_in_A_bool = np.in1d(B_unique_sorted, A, assume_unique=True)

• B_unique_sorted包含排序的B中的唯一值。
• B_idx为这些值保存原始B中的索引。
• B_in_A_bool是一个布尔数组，其大小为B_unique_sorted 存储B_unique_sorted中的值是否位于A。
  注意：我需要在A中查找（B的唯一vals），因为我需要返回有关B_idx
  注意：我假设A已经是唯一的。

现在您可以使用B_in_A_bool来获取公共VAL

B_unique_sorted[B_in_A_bool]

以及它们各自在原始B中的指数

B_idx[B_in_A_bool]

最后，我假设这比纯Python for循环快得多，尽管我没有测试它。