Как передать в функцию массив списков numpy (причудливая индексация)

Question

Я хотел бы использовать эту функцию в векторизованном виде с numpy:

def example(testing, index):
    return np.sum(testing[index], axis = 1)

Предположим, мы создаем наш тестовый массив и массив индексов:

test = np.arange(0, 20)
indices = np.array([[0, 2], [1, 3], [0, 3], [1, 2], [3, 4]])

It делает то, что я ожидаю, суммируя элементы с индексами в списке для каждого списка в массиве:

Input: example(test, indices)
Output: [2 4 3 3 7]

Однако, если я попытаюсь работать со списками индексов переменной длины,

indices = np.array([[0, 2, 3], [1, 3], [0, 3], [1, 2], [3, 4]])
Input: example(test, indices)
Output: IndexError: arrays used as indices must be of integer (or boolean) type

Я не могу заставить numpy выполнить эту функцию без того, чтобы просто перебирать массив indices. Я понимаю, что numpy создает двумерный массив в первом случае, но не во втором, но я не уверен, почему он не может выполнить векторизованную поэлементную операцию с одномерным массивом. Поскольку эти массивы на самом деле довольно большие в реальной жизни и используются в распараллеленной функции, я хотел бы сделать это векторизованным способом в стиле numpy, как в первом примере.

Answer 1

Ваша первая версия indices действительно представляет собой двумерный массив целых чисел - indices.shape возвращает (5, 2), а print(indices) дает:

array([[0, 2],
       [1, 3],
       [0, 3],
       [1, 2],
       [3, 4]])

Вторая, однако, является 1D массив объектов списка :

array([list([0, 2, 3]), list([1, 3]), list([0, 3]), list([1, 2]),
       list([3, 4])], dtype=object)

Здесь просто нечего векторизовать, потому что эти списки не являются numpy векторами . Я считаю, что лучшее, что вы можете сделать для векторизации своих операций, - это создать indices в виде списка массивов, а не массива списков:

indices = [np.array(x) for x in [[0, 2, 3], [1, 3], [0, 3], [1, 2], [3, 4]]]

, а затем переопределить функцию для обработки этого:

def example(testing, index):
    return [ np.sum(testing[i]) for i in index ]

print(example(test, indices)) 
# prints [5, 4, 3, 3, 7]