在numpy（Python）中是否有类似于Matlab中'ismember'的函数？

Question

我正在努力寻找一个Numpy等效的解决方案，用于实现Matlab中的特定编码“模式”，使用ismember。

不幸的是，这段代码往往是我Matlab脚本中耗费大部分时间的地方，因此我想找到一个高效的Numpy等效方案。

基本模式包括将子集映射到更大的网格上。我有一组键值对作为平行数组存储，我想将这些值插入到以相同方式存储的更大的键值对列表中。

具体而言，假设我有季度GDP数据，我将其映射到按月计算的时间网格中，如下所示。

quarters = [200712 200803 200806 200809 200812 200903];
gdp_q = [10.1 10.5 11.1 11.8 10.9 10.3];
months = 200801 : 200812;
gdp_m = NaN(size(months));
[tf, loc] = ismember(quarters, months);
gdp_m(loc(tf)) = gdp_q(tf);

请注意，并非所有季度都出现在月份列表中，因此需要使用tf和loc变量。

我在StackOverflow上看到过类似的问题，但它们要么只提供纯Python解决方案（这里），要么在使用numpy时不返回loc参数（这里）。

在我的特定应用领域中，这种代码模式经常出现，并占据了我函数大部分的CPU时间，因此这里需要一个高效的解决方案。

也欢迎提供评论或重新设计建议。

- snth

---

如果你自己实现的话，可以考虑以下两种方法：1. 对于对象采用哈希函数，既然已有数字，就将其排序并使用二分查找。2. 另一个方法是使用哈希表。 - Mikhail Poda

我认为这个 Alex Martelli 的回答 是你能得到的最好的。 - Sven Marnach

Answer 1

如果months已经排序，使用np.searchsorted。否则，先排序，然后再使用np.searchsorted：

import numpy as np
quarters = np.array([200712, 200803, 200806, 200809, 200812, 200903])
months = np.arange(200801, 200813)
loc = np.searchsorted(months, quarters)

np.searchsorted 返回插入位置。如果存在数据不在正确范围内的可能性，您可能需要之后进行检查：

valid = (quarters <= months.max()) & (quarters >= months.min())
loc = loc[valid]

这是一个O(N log N)的解决方案。如果在运行时方面仍然是你程序中的瓶颈，那么你可以使用哈希方案在C(++)中执行此子例程，这将是O(N)(当然还能避免一些常数因素)。

- luispedro

---

谢谢。你能简要介绍一下哈希方案背后的思想吗？在我的情况下，大多数时候两个关键字数组可能都已排序，因此一个线性O（N）方案，只需并行遍历两个数组即可。不过我总是非常犹豫自己编写任何C代码 - 主要是因为我没有花时间真正调查如何链接自己的扩展。 - snth

我在考虑类似于“hash = dict((val,i) for i,val in enumerate(months)); result = [hash[j] for j in quarters if j in months]”这样的代码，但是用C语言编写。我会使用C++以便能够使用STL哈希类型。 - luispedro

Answer 2

我认为你可以重新设计原始的MATLAB代码示例，使其不使用ISMEMBER函数。这可以加快MATLAB代码的速度，并使其更容易在Python中重新实现（如果你仍然想要）：

quarters = [200712 200803 200806 200809 200812 200903];
gdp_q = [10.1 10.5 11.1 11.8 10.9 10.3];

monthStart = 200801;              %# Starting month value
monthEnd = 200812;                %# Ending month value
nMonths = monthEnd-monthStart+1;  %# Number of months
gdp_m = NaN(1,nMonths);           %# Initialize gdp_m

quarters = quarters-monthStart+1;  %# Shift quarter values so they can be
                                   %#   used as indices into gdp_m
index = (quarters >= 1) & (quarters <= nMonths);  %# Logical index of quarters
                                                  %#   within month range
gdp_m(quarters(index)) = gdp_q(index);  %# Move values from gdp_q to gdp_m

- gnovice

---

+1：ismember函数会执行各种额外的操作，比如调用“unique”函数，这在你的情况下是不必要的，所以你可以简化Matlab（或numpy）代码。 - Jonas

Answer 3

尝试使用pypi中的ismember库。

pip install ismember

例子：

# Import library
from ismember import ismember

# data
quarters = np.array([200712, 200803, 200806, 200809, 200812, 200903])
months = np.arange(200801, 200812)

# Lookup
Iloc,idx=ismember(quarters, months)

# Iloc is boolean defining existence of quarters in months 
print(Iloc)
# [False,  True,  True,  True, False, False]

# index of months that exists in quarters
print(idx)
# [2, 5, 8]

print(months[idx])
[200803, 200806, 200809]

print(quarters[Iloc])
[200803, 200806, 200809]

# These vectors will match
quarters[Iloc]==months[idx]

- erdogant