Equivalente de Matlab 'ismember' en numpy (Python)?

Question

Estoy luchando por encontrar un equivalente de Numpy para un "patrón" de codificación de Matlab particular usando ismember.

Desafortunadamente, este código tiende a ser el que ocupa la mayor parte del tiempo en mis scripts de Matlab, así que quiero encontrar un eficiente Numpy equivalente.

El patrón básico consiste en asignar un subconjunto a una grilla más grande. Tengo un conjunto de pares de valores clave almacenados como matrices paralelas y quiero insertar estos valores en una lista más grande de pares clave de valores almacenados de la misma manera.

Para ser más concretos, decir que tengo datos del PIB trimestrales que asigno a una cuadrícula de tiempo mensual de la siguiente manera.

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); 

Tenga en cuenta que no todos los cuartos aparecen en la lista de meses, por lo tanto la TF y la loc se requieren variables.

He visto preguntas similares sobre StackOverflow pero, o simplemente dar una solución pura de Python (here), o cuando se utiliza numpy entonces no se devuelve el loc argumento (here).

En mi área de aplicación particular, este patrón de código en particular tiende a surgir una y otra vez y consume la mayor parte del tiempo de CPU de mis funciones, por lo que una solución eficiente aquí es realmente crucial para mí.

Comentarios o sugerencias de rediseño también son bienvenidos.

Answer 1

Si se ordenan meses, utilizar np.searchsorted. De lo contrario, ordenar y luego usar 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 devuelve la posición de inserción. Si existe la posibilidad de que sus datos no es ni siquiera en el rango correcto, es posible que desee tener un control después:

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

Este es un O (N log N) solución. Si esto sigue siendo un gran problema en su programa en términos de tiempo de ejecución, puede hacer esta única subrutina en C (++) usando un esquema de hash, que sería O (N) (además de evitar algunos factores constantes, por supuesto).

Answer 2

Creo que puede rediseñar la muestra del código MATLAB original que proporciona para que no utilice la función ISMEMBER. Esto puede acelerar el código de MATLAB y hacerlo más fácil volver a implementar en Python si todavía quiere:

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