я есть две интересующие матрицы, первая - это матрица «слов», с двумя столбцами: идентификатор документа и термин «идентификатор». Например:

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

Кроме того, у меня есть «индексная» матрица, где каждая строка в матрице содержит индекс первой и последней строки для данного идентификатора документа в пакете с матрицей слов. Пример: строка 0 является первым и последним индексом для идентификатора документа 0. Например:

index[0:4]

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

То, что я хотел бы сделать, это взять случайный образец идентификаторов документов и получить всю сумку строк слов для этих идентификаторов документов. Матрица пакета слов составляет примерно 150M строк (~ 1,5 ГБ), поэтому использование numpy.in1d ​​() слишком медленное. Нам нужно быстро вернуть их для подачи в последующую задачу.

Наивное решение, которое я нашел, заключается в следующем:

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

Общий образец

Типичным примером для постановки проблемы будет что-то вроде этого -

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

Ожидаемый результат будет -

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