Tengo dos matrices de interés, la primera es una matriz de "bolsa de palabras", con dos columnas: la identificación del documento y la identificación del término. Por ejemplo:

bow[0:10]

Out[1]:
    array([[ 0, 10],
           [ 0, 12],
           [ 0, 19],
           [ 0, 20],
           [ 1,  9],
           [ 1, 24],
           [ 2, 33],
           [ 2, 34],
           [ 2, 35],
           [ 3, 2]])

Además, tengo una matriz de "índice", donde cada fila de la matriz contiene el índice de la primera y la última fila para una ID de documento dada en la matriz de bolsa de palabras. Ej: la fila 0 es el primer y último índice para la identificación de documento 0. Por ejemplo:

index[0:4]

Out[2]:
    array([[ 0,  4],
           [ 4,  6],
           [ 6,  9],
           [ 9, 10]])

Lo que me gustaría hacer es tomar una muestra aleatoria de ID de documentos y obtener toda la bolsa de filas de palabras para esas ID de documentos. La matriz de bolsa de palabras tiene aproximadamente 150 millones de filas (~ 1.5 Gb), por lo que usar numpy.in1d () es demasiado lento. Necesitamos devolverlos rápidamente para alimentar una tarea aguas abajo.

La solución ingenua que se me ocurrió es la siguiente:

def get_rows(ids):
    indices = np.concatenate([np.arange(x1, x2) for x1,x2 in index[ids]])
    return bow[indices]

get_rows([4,10,3,5])

Muestra genérica

Una muestra genérica para plantear el problema sería con algo como esto:

indices = np.array([[ 4, 7],
                    [10,16],
                    [11,18]]

El resultado esperado sería:

array([ 4,  5,  6, 10, 11, 12, 13, 14, 15, 11, 12, 13, 14, 15, 16, 17])