CODEWITHSUNDEEP
pythonnumpyvectorization
lumpy
lumpy

Reputation: 87

How to vectorise these two nested for-loops?

I have the following code:

import numpy as np

my_array = np.zeros((42, 123, 2021))

assert another_array.shape == (42, 123)

for i in range(42):
    for j in range(123):
        my_array[i, j, another_array[i, j]] = 1

where it is assumed that the values of another_array stay inside the correct range (i.e. the values of another_array are integers between 0 and 2020).

I would like to get rid of the two for-loops. Is there a way to vectorise something like this?

Upvotes: 3

Views: 74

Answers (2)

hpaulj
hpaulj

Reputation: 231335

Try:

my_array = np.zeros((42, 123, 2021))
# assert another_array.shape == (42, 123)
my_array[ np.arange(42)[:,None], np.arange(123), another_array] = 1

The idea is to replace the i,j with a pair of ranges that broadcast to (42,123) to match the 3rd axis index array.

Upvotes: 1

lumpy
lumpy

Reputation: 87

I got it:

import numpy as np

my_array = np.zeros((42, 123, 2021))
my_array = my_array.reshape(-1, 2021)

assert another_array.shape == (42, 123)
another_array = another_array.flatten()

my_array[range(my_array.shape[0]), another_array] = 1
my_array = my_array.reshape(42, 123, 2021)

Upvotes: 1

Related Questions