Vectorized Update numpy array using another numpy array elements as index

Question

Let A,C and B be numpy arrays with the same number of rows. I want to update 0th element of A[0], 2nd element of A[1] etc. That is, update B[i]th element of A[i] to C[i]

import numpy as np
A = np.array([[1,2,3],[3,4,5],[5,6,7],[0,8,9],[3,7,5]])
B = np.array([0,2,1,2,0])
C = np.array([8,9,6,5,4])
for i in range(5):
A[i, B[i]] = C[i]
print ("FOR", A)
A = np.array([[1,2,3],[3,4,5],[5,6,7],[0,8,9],[3,7,5]])
A[:,B[:]] = C[:]
print ("Vectorized, A", A)

Output:

FOR [[8 2 3]
[3 4 9]
[5 6 7]
[0 8 5]
[4 7 5]]
Vectorized, A [[4 6 5]
[4 6 5]
[4 6 5]
[4 6 5]
[4 6 5]]

The for loop and vectorization gave different results. I am unsure how to vectorize this for loop using Numpy.

The description in the text makes it sound like the last row should be updated. I assumed range(4) is a typographical error. — Warren Weckesser
– Warren Weckesser, Commented Jun 27, 2018 at 18:25
@WarrenWeckesser Yeah, that's what I assumed at first and hesitated to post an answer but still as I explained in my answer it's simply possible to do so. — Kasravnd
– Kasravnd, Commented Jun 27, 2018 at 18:28

Kasravnd · Accepted Answer · 2018-06-27 18:24:55Z

1

The reason that your approach doesn't work is that you're passing the whole B as the column index and replace them with C instead you need to specify both row index and column index. Since you just want to change the first 4 rows you can simply use np.arange(4) to select the rows B[:4] the columns and C[:4] the replacement items.

In [26]: A[np.arange(4),B[:4]] = C[:4]

In [27]: A
Out[27]: 
array([[8, 2, 3],
       [3, 4, 9],
       [5, 6, 7],
       [0, 8, 5],
       [3, 7, 5]])

Note that if you wanna update the whole array, as mentioned in comments by @Warren you can use following approach:

A[np.arange(A.shape[0]), B] = C

answered Jun 27, 2018 at 18:24

Kasravnd

108k19 gold badges167 silver badges195 bronze badges

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

Collectives™ on Stack Overflow

Vectorized Update numpy array using another numpy array elements as index

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related