Есть ли более быстрая реализация следующего кода?

Question

У меня есть одномерный массив numpy, который довольно большой по размеру.Для каждой записи массива мне нужно создать линейно распределенный подмассив до этого значения записи.Вот что я привожу в качестве примера.

import numpy as np
a = np.array([2, 3])
b = np.array([np.linspace(0, i, 4) for i in a])

В этом случае есть линейное пространство размера 4. Последнее утверждение в приведенном выше коде включает for loop, что довольно медленно, если aочень большой.Есть ли хитрость для реализации этого в самой numpy?

Answer 1

Вы можете сформулировать это как внешний продукт :

In [37]: a = np.arange(100000)

In [38]: %timeit np.array([np.linspace(0, i, 4) for i in a])
1 loop, best of 3: 1.3 s per loop

In [39]: %timeit np.outer(a, np.linspace(0, 1, 4))
1000 loops, best of 3: 1.44 ms per loop

Идея состоит в том, чтобы взять единицу linspace и затем масштабировать ее отдельно по каждому элементу a.

Как вы можете видеть, это дает увеличение скорости на ~ 1000x для n=100000.

Для полноты я упомяну, что этот код имеет немного отличающиеся свойства округления, чем ваша исходная версия (вероятно,не проблема в практических приложениях):

In [52]: np.max(np.abs(np.array([np.linspace(0, i, 4) for i in a]) -
    ...:               np.outer(a, np.linspace(0, 1, 4))))
Out[52]: 1.4551915228366852e-11

PS Альтернативный способ выразить идею заключается в использовании поэлементного умножения с вещанием (на основе предложения @Scott Gigante):

In [55]: %timeit a[:, np.newaxis] * np.linspace(0, 1, 4)
1000 loops, best of 3: 1.48 ms per loop

PPS См. Комментарии ниже, чтобы узнать, как сделать это быстрее.