Эквивалент Matlab 'ismember' в numpy (Python)?

Я изо всех сил пытаюсь найти эквивалент Numpy для конкретного шаблона Matlab, используя ismember.

К сожалению, этот код, как правило, больше всего времени тратится на мои скрипты Matlab, поэтому я хочу найти эффективный эквивалент Numpy.

Основной шаблон состоит в отображении подмножества на большую сетку. У меня есть набор пар значений ключей, хранящихся как параллельные массивы, и я хочу вставить эти значения в более широкий список пар значений ключа, сохраненных таким же образом.

Для конкретности говорят, что у меня есть данные по квартальному ВВП, которые я накладываю на ежемесячную сетку времени следующим образом.

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 не возвращается (здесь).

В моей конкретной области приложения этот конкретный шаблон кода имеет тенденцию возникать снова и снова и использует большую часть времени процессора для моих функций, поэтому эффективное решение здесь действительно важно для меня.

Комментарии и рекомендации по редизайну также приветствуются.

Ответ 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) (а также избежать некоторых постоянных факторов, конечно).

Ответ 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