1

I want to improve a code where I have to sum over different indices. This is the code:

neighbors = [[1,2,4],[2,0],[1,3],[],[0]]
omega = np.random((5,5,3,3))

#Sum over indices 0 (filtered) and 2 
suma = np.zeros((5,3))
for j in range(5):
    for l in range(3):
        suma[j,l] += np.sum(omega[neighbors[j],j,:,l])

Now I improved a little bit my code using the axis parameter of sum:

suma2 = np.zeros((5,3))
for j in range(5):
    suma2[j,:] += np.sum(omega[neighbors[j],j,:,:],axis=(0,1))

I want to know how I can avoid the first loop. I tried creating an array of booleans:

neighbors_bool = np.full((5,5),False,dtype=bool)
for j in range(5):
    neighbors_bool[neighbors[j],j] = True

But I don't know how to put it in the sum.

oscarcapote
  • 417
  • 1
  • 4
  • 16
  • The fact that the sublists of `neighbors` vary in size from 0 to 3 is a pretty good indicator that each sublist will have to be handled on its own. See this recent question for more details: [Apply function along ranges in numpy array](https://stackoverflow.com/questions/54540649/apply-function-along-ranges-in-numpy-array) – hpaulj Feb 07 '19 at 17:58

0 Answers0