Chen
Chen

Reputation: 81

Numpy: Replacing part of array by another entire array based on selected index

I have two np array A and B and one index list, where:

A = np.arange(12).reshape(3,2,2)
>>> A
array([[[ 0,  1],
        [ 2,  3]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 8,  9],
        [10, 11]]])

>>> B:
array([[[10, 10],
        [10, 10]],

       [[20, 20],
        [20, 20]]])

index_list = [0,2]

I would like to replace array A by entire array B based on index the index_list gives, which is 0 and 2 (corresponding to A's index at axis=0):

# desired output:
array([[[ 10,  10],
        [ 10,  10]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 20,  20],
        [20, 20]]])

I have been thinking about implementing this process in a more efficient way, but only come up with two loops:

row_count = 0
for i,e in enumerate(A):
    for idx in index_list:
        if i == idx:
            A[i,:,:] = B[row_count,:,:]
            row_count += 1

which is computational expensive because I have two huge np arrays to process... it would be very appreciated if anyone have any idea of how to implement this process in vectorization way or in a more efficient way. Thanks!

Upvotes: 0

Views: 796

Answers (1)

Amir Charkhi
Amir Charkhi

Reputation: 836

Have you tried just assigning like this:

A[index_list[0]] = B[0]
A[index_list[1]] = B[1]

I guess if you had a bigger number of indexes in the index_list, and more to replace, you could then make a loop.

Upvotes: 2

Related Questions