Группировка элементов массива NumPy по сумме индексов

Question

У меня есть несколько больших numpy массивов измерений 30 * 30 * 30, по которым мне нужно пройти массив, получить сумму каждого индексного триплета и скомпоновать эти элементы на эту сумму. Например, рассмотрим этот простой массив 2 * 2:

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

Этот массив имеет индексы [0,0], [0,1], [1,0] и [1,1]. Эта процедура возвращает список: [2, [3,0], 1], потому что 2 в массиве test имеет индексную сумму 0, 3 и 0 имеют индексную сумму 1, а 1 имеет индексную сумму 2. I Я знаю, что метод грубой силы итерации массива NumPy и проверки суммы будет работать, но он слишком неэффективен для моего реального случая с большим N (= 30) и несколькими массивами. Любые входы по использованию подпрограмм NumPy для выполнения sh этой группировки будут оценены. Заранее спасибо.

Answer 1

Вот один способ, который должен быть достаточно быстрым, но не сверхбыстрым: 30x30x30 занимает 20 ms на моей машине.

import numpy as np

# make example
dims = 2,3,4
a = np.arange(np.prod(dims),0,-1).reshape(dims)

# create and sort indices
idx = sum(np.ogrid[tuple(map(slice,dims))])
srt = idx.ravel().argsort(kind='stable')

# use order to arrange and split data
asrt = a.ravel()[srt]
spltpts = idx.ravel().searchsorted(np.arange(1,np.sum(dims)-len(dims)+1),sorter=srt)
out = np.split(asrt,spltpts)

# admire
out
# [array([24]), array([23, 20, 12]), array([22, 19, 16, 11,  8]), array([21, 18, 15, 10,  7,  4]), array([17, 14,  9,  6,  3]), array([13,  5,  2]), array([1])]

Answer 2

Вы можете процедурно создать список индексных кортежей и использовать его, но, возможно, вы получаете кодовую константу, которая слишком велика, чтобы быть эффективной. [(0,0), [(1,0), (0,1)], (1,1)],

Так что вам нужна функция для генерации этих индексов на лету для n- массив измерений.

Для одного измерения, тривиальный счет / приращение

   [(0),(1),(2),...] 

Для второго, используйте стратегию одного измерения для первого измерения, уменьшите первое и увеличьте второе до заполните.

   [(0...)...,(1...)...,(2...)...,...] 
   [[(0,0)],[(1,0),(0,1)],[(2,0),(1,1),(0,2)],[...],...]

Обратите внимание, что некоторые из них будут находиться за пределами массива примеров. Ваш генератор должен будет включать проверку границ.

Затем три измерения, дать первым двум измерениям обработку как указано выше, но в конце уменьшите первое измерение, увеличьте третье, повторяйте до тех пор, пока все не будет сделано Общий алгоритм заключается в том, что вы должны go генерировать индексы на лету, а не два больших массива конкурировать за кеш и ввод-вывод.

Создание python или его эквивалента оставлено пользователю в качестве упражнения.