一尘不染

Python等同于MATLAB的“ ismember”函数

python

经过多次尝试优化代码之后,似乎最后的资源就是尝试使用多个内核来运行以下代码。我不确切地知道如何转换/重组我的代码,以便它可以使用多个内核更快地运行。如果能得到指导以实现最终目标,我将不胜感激。最终目标是能够对数组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,对于A中属于B的每个值,在B中包含最低的索引。无论A是否不是B的成员,输出数组Locb都包含0。

主要问题之一是我需要能够尽可能高效地执行此操作。为了测试,我有两个700k元素的数组。创建一个生成器并检查生成器的值似乎并不能很快完成工作。


阅读 374

收藏
2020-12-20

共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

您最初的实现需要对B中的元素对A中的每个元素进行全面扫描,使其成为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)表达式中的第二个参数更改为要返回的值。

2020-12-20