Блог

Главная - Вопросы по программированию - Векторизованный индекс массивов

Векторизованный индекс массивов

#python #numpy

#python #numpy

Вопрос:

Изначально у меня было что-то вроде этого:

 a = 1   # Some randomly generated positive integer
b = -1  # Some randomly generated negative integer
c = 0   # Constant 0

i = 0  # Randomly picked from (0, 1, 2)

d = [a, b, c][i]

Я хотел бы векторизовать это так, чтобы можно было генерировать много выборок

Итак, у меня есть три массива длиной N, индексный массив длиной N, и я хотел бы использовать этот индексный массив для выбора одного из трех массивов

 a = np.array([1, 2, 3, 4])
b = np.array([-1, -2, -3, -4])
c = np.array([0, 0, 0, 0])

i = np.array([2, 1, 2, 0])

d = np.array([a, b, c])[i]  # Doesn't work
# Would like the result:
d = np.array([0, -2, 0, 4])

d = a * (i == 0) b * (i == 1) c * (i == 2) работает, но наверняка есть способ, который больше похож на невекторизованный код

Ответ №1:

Создайте 2-d массив из трех массивов, затем используйте целочисленное индексирование

 >>> e = np.vstack([a,b,c])
>>> i = np.array([2, 1, 2, 0])
>>> e[(i,np.arange(i.shape[0]))]
array([ 0, -2,  0,  4])
>>>

Комментарии:

1. В моем исходном коде есть size параметр, который может быть None для невекторизованной версии, на которую это можно легко расширить np.vstack([a, b, c])[(i, 0 if size is None else np.arange(size))] . Спасибо! (Хотя это не распространяется на size=[shape_0, shape_1] )

2. @frankden — size=[shape_0, shape_1] — в результате получается три 2-d массива?

Ответ №2:

Обратите внимание, что ваш ответ находится по диагонали

  np.array([a, b, c])[i]

итак, вы можете перейти:

  np.array([a, b, c])[i].diagonal()
Метки: