Python相当于MATLAB的"ismember"函数

Question

在尝试优化代码的许多尝试之后,似乎最后一个资源是尝试使用多个核运行下面的代码.我不知道如何转换/重新构建我的代码,以便使用多个内核可以更快地运行.如果我能获得指导以实现最终目标,我将不胜感激.最终目标是能够尽快为阵列A和B运行此代码,其中每个阵列包含大约700,000个元素.这是使用小数组的代码.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部分.

来自Matlab:Locb,包含B中每个值的最低索引,A是B的成员.输出数组Locb在A不是B的成员的地方包含0

其中一个主要问题是我需要能够尽可能高效地执行此操作.为了测试,我有两个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进行一次全扫描才能生成字典Bset.通过使用dict,您可以有效地使A中每个元素的每个元素的查找为常量,从而进行操作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

sfstewman的优秀答案很可能为您解决了这个问题.

我只想补充一下如何在numpy中完全相同.

我使用了numpy 独特的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的唯一值),因为我需要输出相对于B_idx
  注意:我认为它A已经是唯一的.

现在你可以B_in_A_bool用来获得常见的val

B_unique_sorted[B_in_A_bool]

和他们各自的原始指数 B

B_idx[B_in_A_bool]

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