Modifying multidimensional array at a specific index in Python

Question

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.

You write about "array" and "shape" and use [:,:,2] in your code. Are you using numpy arrays? Do you want your answer to be limited to regular Python or should it use numpy? By the way, your code example is elegant, since it is clear. Do you really need something faster? — Rory Daulton
– Rory Daulton, Commented Feb 4, 2018 at 0:47
@RoryDaulton I am using numpy arrays. I will edit that, sorry. I was thinking that it would be possible to do it with less lines of code and maybe faster using vectorization. — Molina12
– Molina12, Commented Feb 4, 2018 at 0:54
@Molina12, i'm not sure it'll help in this case but consider using numba.jit(nopython=True) on your for loop function. — jpp
– jpp, Commented Feb 4, 2018 at 0:59

Matt Hall · Accepted Answer · 2018-02-04 01:02:50Z

2

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.

answered Feb 4, 2018 at 1:02

Matt Hall

8,1691 gold badge28 silver badges38 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Modifying multidimensional array at a specific index in Python

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related