Эффективный способ l oop через ортодиагональные индексы в порядке

Question

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

import numpy as np

len_x, len_y = 50, 50 #they don't have to equal
index_arr = np.add.outer(np.arange(len_x), np.arange(len_y))

В настоящее время я перебираю вот так:

for i in range(np.max(index_arr)):
    orthodiag_indices = zip(*np.where(index_arr == i))
        for index in orthodiag_indices:
            # DO FUNCTION OF index #

У меня есть произвольная функция кортежа индекса, index и других параметров вне этого l oop. Такое ощущение, что мне не нужен второй for l oop, и я должен быть в состоянии сделать все это за один l oop. Кроме того, я делаю много вызовов функций из zip(*np.where(index_arr == i)) для каждого i. Какой самый эффективный способ сделать это?

Редактировать: следует упомянуть, что важно, чтобы функция применялась к index_arr == i по порядку, то есть сначала она делала 0, затем 1, затем 2 и c. (порядок второго l oop не имеет значения).

Редактировать 2: Я думаю, что мне нужен способ получить индексы [(0,0), (0,1), (1,0), (2,0), (1,1), (2,0), ...] эффективно. Я не думаю, что смогу применить векторизованную функцию, потому что я заполняю массив np.zeros((len_x, len_y)) и возвращаюсь к первому редактированию, порядок имеет значение.

Answer 1

Я думаю, что itertools.product () будет здесь использоваться

import itertools as it
x,y = 2,3
a=list(it.product(range(x),range(y))

, что дает a как

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

Если они вам нужны в затем,

b=np.argsort(np.sum(a,1))
np.array(a)[b]

, что дает,

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

Надеюсь, что поможет!

Answer 2

Вы можете использовать tril/triu_indices. Поскольку порядок (бывшего) внутреннего l oop не имеет значения, размеры могут быть заменены при необходимости, я буду считать L>=S:

L,S = 4,3

a0,a1 = np.tril_indices(L,0,S)
b0,b1 = np.triu_indices(S,1)
C0 = np.concatenate([a0-a1,b0+L-b1])
C1 = np.concatenate([a1,b1])
*zip(C0,C1),
# ((0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2), (3, 0), (2, 1), (1, 2), (3, 1), (2, 2), (3, 2))