Как да разопаковате кортеж за цикъл, без да сте специфични за измерението

Бих искал да направя нещо подобно:

    if dim==2:
        a,b=grid_shape
        for i in range(a):
            for j in range(b):
                A[i,j] = ...things...

където dim е просто броят на елементите в моя кортеж grid_shape. A е numpy масив с измерение dim. Има ли начин да го направите, без да сте специфични за измерението? Без да се налага да пишете грозен код като

    if dim==2:
        a,b=grid_shape
        for i in range(a):
            for j in range(b):
                A[i,j] = ...things...
    if dim==3:
        a,b,c=grid_shape
        for i in range(a):
            for j in range(b):
                for k in range(c):
                    A[i,j,k] = ...things...

person LukeMathWalker    schedule 16.10.2016    source източник
comment
IIUC, търсите ndenumerate. Можеш ли да потвърдиш?   -  person DSM    schedule 16.10.2016
comment
Не точно, въпреки че тази функция ще се окаже полезна в друга част от рутината, така че благодаря все пак :D   -  person LukeMathWalker    schedule 16.10.2016


Отговори (2)


Използвайки itertools, можете да го направите по следния начин:

for index in itertools.product(*(range(x) for x in grid_shape)):
    A[index] = ...things...

Това разчита на няколко трика. Първо, itertools.product() е функция, която генерира кортежи от итерируеми.

for i in range(a):
    for j in range(b):
        index = i,j
        do_something_with(index)

може да се сведе до

for index in itertools.product(range(a),range(b)):
    do_something_with(index)

Това работи за произволен брой аргументи на itertools.product(), така че можете ефективно да създавате вложени цикли с произволна дълбочина.

Другият трик е да преобразувате формата на вашата мрежа в аргументите за itertools.product:

(range(x) for x in grid_shape)

е еквивалентно на

(range(grid_shape[0]),range(grid_shape[1]),...)

Тоест, това е набор от диапазони за всяко измерение grid_shape. Използването на * след това разширява това в аргументите.

itertools.product(*(range(x1),range(x2),...))

е еквивалентно на

itertools.product(range(x1),range(x2),...)

Освен това, тъй като A[i,j,k] е еквивалентно на A[(i,j,k)], можем просто да използваме A[index] директно.

Както посочва DSM, тъй като използвате numpy, можете да намалите

itertools.product(*(for range(x) for x in grid_shape))

to

numpy.ndindex(grid_shape)

Така става последният цикъл

for index in numpy.ndindex(grid_shape):
    A[index] = ...things...
person Vaughn Cato    schedule 16.10.2016
comment
Това изглежда работи! Бихте ли ми обяснили машината зад това? - person LukeMathWalker; 16.10.2016
comment
Вижте np.ndindex, ако всичко, което искаме, е индексът, а не стойността. - person DSM; 16.10.2016
comment
И двамата бяхте изключително полезни! Благодаря ви :D - person LukeMathWalker; 16.10.2016

Можете да хванете останалата част от кортежа, като поставите звезда пред последната променлива и направите масив, като поставите скоби около него.

>>> tupl = ((1, 2), 3, 4, 5, 6)
>>> a, *b = tupl
>>> a
(1, 2)
>>> b
[3, 4, 5, 6]
>>> 

И тогава можете да преминете през b. Така че би изглеждало нещо подобно

a,*b=grid_shape
for i in a:
    for j in range(i):
        for k in b:
            for l in range(k):
                A[j, l] = ...things...
person rassar    schedule 16.10.2016
comment
Това няма да работи, защото A ще се окаже двумерен масив, дори ако b има повече от един елемент. - person LukeMathWalker; 16.10.2016
comment
това оправя ли го? - person rassar; 16.10.2016