Перебор массива с сохранением индексов

Question

Рассмотрим следующий фрагмент кода:

n_samples, n_rows, n_cols, n_boxes_per_cell, _ = Y_pred.shape
    for example in range(n_samples):
        for x_grid in range(n_rows):
            for y_grid in range(n_cols):
                for bnd_box in range(n_boxes_per_cell):
                    bnd_box_label = Y_pred[example, x_grid, y_grid, bnd_box]
                    do_some_stuff(bnd_box_label, x_grid, y_grid)

Как получить функционально эквивалентный код с не более чем одной явной итерацией?Обратите внимание, что мне нужны индексы x_grid и y_grid.

Answer 1

Я не уверен, что это то, что вы ищете, но вы всегда можете собрать генератор из нескольких итераций:

all_combinations = ((a, b, c, d) for a in range(n_samples)
                                 for b in range(n_rows) 
                                 for c in range(n_cols) 
                                 for d in range(n_boxes_per_cell))

for examples, x_grid, y_grid, bnd_box in all_combinations:
    do stuff

Это то же самое, что и использование itertools.product (*итерируемые) и действительны для любой итерируемой, а не только итерации по индексам / целым числам.

Answer 2

Вы можете использовать np.ndindex:

for example, x_grid, y_grid, bnd_box in np.ndindex(Y_pred.shape[:4]):
    etc.