CODEWITHSUNDEEP
pythonpython-3.xnumpy
Molina12
Molina12

Reputation: 200

Modifying multidimensional array at a specific index in Python

Let's say I have a numpy array A with shape: (10, 20, 3) I want to modify index 2 of the third array for all the rows and columns. I tried doing it this way:

A[:,:,2] += 100

This partly works but when I want to check if the value of each row and column at index 2 is higher than 200 I cannot unless I use two nested for loops. This works, but it is not fast at all.

This is what I have right now:

for i in range(0,10):
    for j in range(20):
        if (A[i,j,2] + 100) > 200 :
            A[i,j,2] = 200
        else:
            A[i,j,2] += 100

I am looking for a more efficient and elegant way to accomplish this.

Upvotes: 1

Views: 511

Answers (1)

Matt Hall
Matt Hall

Reputation: 8112

I understand you want to add 100 to all the elements in the 2nd index of the 3rd dimension... but if the sum goes over 200, just make it 200.

A rule of thumb in NumPy is you ~never need loops. I think you can achieve what you want directly with indexing, eg:

>>> import numpy as np
>>> A = np.random.random((10, 20, 3))*200
>>> A[...,2] += 100
>>> A[...,2][A[...,2] > 200] = 200
array([[[ 131.94487476,   57.20338449,  200.        ],
        [  22.7506817 ,  145.74740259,  200.        ],
        [  15.44748999,  180.55442849,  189.48182561],
        [ 160.35622709,    2.37465206,  183.96627351],
        [  67.79103218,   70.28339315,  108.90041475],
        [ 129.33682572,  102.6764913 ,  200.        ],
        .
        .
        .

Where A[...,2] is just another way of saying the second index (third element) in the last dimension.

Upvotes: 2

Related Questions