Search code examples
pythonarraysnumpyin-place

Increment elements of a subset of elements in a 3d numpy array python


I have a 3d list formed by

myArray = np.array([[[0]*n for i in range(m)] for j in range(o)])

I have a loop that runs over all elements, and increments the value stored in the current element and a number of elements in the neighborhood of the current element:

myArray[xa:xb][ya:yb][za:zb] += 1.

where xa,xb, etc. are generated according to the current element considered in the loop, and not necessarily the same. In other words, I'd like to increment the values of a given sub-triangle in the 3D list.

However, when I try to address myArray[xa:xb][0][0], I get a list with length that is larger than len(myArray[0]). Not to mention myArray[xa:xb][ya:yb][za:zb] += 1 results in more elements to be incremented by 1 than desired.

I could achieve this by using three nested for loops:

for i in range(xa,xb+1):
    for j in range(ya,yb+1):
        for k in range(za,zb+1):
            myArray[i][j][k] += 1

but this slows down the code a lot. What can I do to achieve this without such a loss of performance?


Solution

  • You were on the right path from the beginning. The following seems to work:

    myArray=np.zeros((o,m,n))
    myArray[xa:xb+1,ya:yb+1,za:zb+1]+=1
    

    Note that index slicing in arrays uses the same boundaries as range in your for loop, thus you have to +1 your end index. The procedure above replicates your triple for loops results, at a fraction of time.