Numpy способ шага вперед для цикла - PullRequest
Новая
       82

Numpy способ шага вперед для цикла

1 голос
/ 04 августа 2020

Я хочу избежать for-l oop с шагом в следующем коде и заменить его кодом Numpy, чтобы ускорить процесс: 

import numpy as np

A = np.array([[0,0,7,0,0,0], [0,0,0,0,5,0]])

for i in range(A.shape[0]):
    for j in range(A.shape[1]-1):
        if A[i,j]==7 and A[i,j+1]==0:
            A[i,j+1]=7

Я знаю, как это сделать с помощью for -l oop без шага. Скажем, A,B,C - это 2D-массивы одинакового размера, тогда этот медленный код: 

for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        if A[i,j]==7 and B[i,j]==0:
            C[i,j]=7

... может быть быстрее с помощью следующей единственной строки numpy кода: 

C[(A==7) & (B==0)]=7

Думаю, должно быть что-то похожее, включая функции np.where и np.roll? Благодарю за вашу помощь!

1 Ответ

1 голос
/ 04 августа 2020

Я почти уверен, что есть способы получше, но если вы его не нашли, вот более быстрый способ: 

for i in range(A.shape[1]-1):
  A = np.where((A==0)&(np.pad(A,((0,0),(1,0)))[:,:-1]==7),7,A)

вывод: 

[[0 0 7 7 7 7]
 [0 0 0 0 5 0]]
...