在numpy(Python)中Matlab'ismember'的等价物？

Question

我正在努力为使用ismember的特定Matlab编码"模式"找到Numpy等价物.

不幸的是,这段代码往往是我的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时间,因此这里的有效解决方案对我来说至关重要.

我们也欢迎提出意见或重新设计建议.

Answer 1

如果按月分类,请使用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)(以及避免一些常数因素,当然).