Python 等效于 MATLAB 的“ismember”函数

Posted 2023-02-23

【中文标题】Python 等效于 MATLAB 的“ismember”函数

【英文标题】：Python equivalent of MATLAB's "ismember" function

【发布时间】：2013-03-29 15:33:46

【问题描述】：

经过多次尝试优化代码后，最后一个资源似乎是尝试使用多个内核运行下面的代码。我不确切知道如何转换/重新构造我的代码，以便它可以使用多个内核运行得更快。如果我能得到指导以实现最终目标，我将不胜感激。最终目标是能够尽可能快地为数组 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 中的最低索引。输出数组 Locb 包含 0，只要 A 不是 B 的成员

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

【参考方案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，您可以有效地将 B 中的每个元素查找为 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) 表达式中的第二个参数更改为您想要返回的值。

【讨论】：

哇，这速度真快！你不知道我多么欣赏你的解决方案。非常感谢 ！您是否使用特定工具来输出性能配置文件？

@z5151 不，这是简单的算法分析。使用Big-O notation：np.where 必须对B 进行线性扫描，这需要O(len(B)) 操作。然后，您使用需要 O(len(A)) 操作的外部循环，使您的原始算法大致为 O(len(A)*len(B)) 操作。生成Bind 需要len(B) 操作。字典实现为hash tables，有常量O(1)查找，所以扫描A为O(len(A))；整体复杂度为O(len(A)+len(B))。

知道了。感谢您的***参考。

@EOL 不，你破坏了密码。返回的元素现在是列表中的最后一个，而不是第一个。我没有在原始代码中使用字典理解是有原因的。

@EOL 据我所知，您可以通过遍历反向范围来使用字典理解： B[i] : i for i in xrange(len(B)-1,-1,-1) 。或者使用反向迭代器： elt : len(B)-i-1) for (i,elt) in enumerate(reversed(B)) 两者都不漂亮（或简单）。第一个假设 B 是可转位的，第二个假设 B 是可逆的。他们还假设随机索引/反向迭代很便宜。如果 B 是一个非常大的链表，那么你的表现将会很糟糕。有没有一种使用推导式来检索仅假设迭代的第一个索引的方法？

【参考方案2】：

sfstewman 的出色回答很可能为您解决了问题。

我想补充一下如何在 numpy 中实现同样的功能。

我使用了 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 的唯一 val），因为我需要要返回的输出与 B_idx 相关

注意：

我假设 A 已经是唯一的。

现在您可以使用 B_in_A_bool 来获取常用 vals

B_unique_sorted[B_in_A_bool]

以及它们各自在原始B中的索引

B_idx[B_in_A_bool]

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

【讨论】：

+1 尽可能使用 numpy，可以通过这种方式实现主要的加速（因为我已经学会了艰难的方式>_

当心！这不会保留索引的顺序！尝试使用 range(1,6) 和 [5,1]。如果不需要索引的顺序，我认为您可以使用 np.in1d() 然后使用 np.nonzero()[0]

在此处查看答案：***.com/questions/33678543/… 以正确顺序获取索引

【参考方案3】：

试试ismember 库。

pip install ismember

简单示例：

# Import library
from ismember import ismember
import numpy as np

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

# Lookup
Iloc,idx = ismember(A, B)
 
# Iloc is boolean defining existence of d in d_unique
print(Iloc)
# [ True False False  True  True]

# indexes of d_unique that exists in d
print(idx)
# [4 4 3]

print(B[idx])
# [3 3 6]

print(A[Iloc])
# [3 3 6]

# These vectors will match
A[Iloc]==B[idx]

速度检查：

from ismember import ismember
from datetime import datetime

t1=[]
t2=[]
# Create some random vectors
ns = np.random.randint(10,10000,1000)

for n in ns:
    a_vec = np.random.randint(0,100,n)
    b_vec = np.random.randint(0,100,n)

    # Run stack version
    start = datetime.now()
    out1=ismember_stack(a_vec, b_vec)
    end = datetime.now()
    t1.append(end - start)

    # Run ismember
    start = datetime.now()
    out2=ismember(a_vec, b_vec)
    end = datetime.now()
    t2.append(end - start)


print(np.sum(t1))
# 0:00:07.778331

print(np.sum(t2))
# 0:00:04.609801

# %%
def ismember_stack(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

pypi 的 ismember 函数快了近 2 倍。

大向量，例如 700000 个元素：

from ismember import ismember
from datetime import datetime

A = np.random.randint(0,100,700000)
B = np.random.randint(0,100,700000)

# Lookup
start = datetime.now()
Iloc,idx = ismember(A, B)
end = datetime.now()

# Print time
print(end-start)
# 0:00:01.194801

【讨论】：

@EOL 你觉得这个解决方案怎么样？

【参考方案4】：

尝试使用列表推导；

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]

通常，列表推导比 for 循环快得多。

获得等长列表；

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() 来加快操作速度。

【讨论】：

输出数组需要保持相同大小。

@z5151：查看增强的答案。如果需要，您可以将 lambda 表达式更改为返回 0 而不是 False，但这会掩盖结果中的真实 0。

这对于元素数量很少的数组很有用。感谢您强调列表推导比循环快得多。

您的答案返回的是元素，而不是 B 中元素的索引。

【参考方案5】：

这是返回与 MATLAB 匹配的输出参数 [Lia, Locb] 的确切 MATLAB 等效项，但 Python 0 也是有效索引。因此，此函数不返回 0。它本质上返回 Locb(Locb>0)。性能也与MATLAB相当。

def ismember(a_vec, b_vec):
    """ MATLAB equivalent ismember function """

    bool_ind = np.isin(a_vec,b_vec)
    common = a[bool_ind]
    common_unique, common_inv  = np.unique(common, return_inverse=True)     # common = common_unique[common_inv]
    b_unique, b_ind = np.unique(b_vec, return_index=True)  # b_unique = b_vec[b_ind]
    common_ind = b_ind[np.isin(b_unique, common_unique, assume_unique=True)]
    return bool_ind, common_ind[common_inv]

这里是一个稍慢（~5x）但不使用独特功能的替代实现：

def ismember(a_vec, b_vec):
    ''' MATLAB equivalent ismember function. Slower than above implementation'''
    b_dict = b_vec[i]: i for i in range(0, len(b_vec))
    indices = [b_dict.get(x) for x in a_vec if b_dict.get(x) is not None]
    booleans = np.in1d(a_vec, b_vec)
    return booleans, np.array(indices, dtype=int)