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

Question

Я изо всех сил пытаюсь найти эквивалент 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 isn 't вернул ( здесь ).

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

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

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

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