CODEWITHSUNDEEP
pythonarraysmatlab
vzografos
vzografos

Reputation: 117

Slicing a 2d array in python via another array in one line

Apologies is the title is not correct. I didn't know how to describe exactly what I am looking for. So coming from Matlab I want to be able to do the following in Python in one line if possible

Say I have an index array:

index_array= np.array([0,1,1,0,0,0,1])

and a data array: data_array = np.zeros([len(index_array),2])

I want to place a value (e.g. 100) where index_array=0 to the data_array[:,0] and where index_array=1 to data_array[:,1]

I matlab you could do it with one line. Something like data_array(index_array)=100

The best I could figure out in python is this

data_array [index_array==0,0]=100

data_array [index_array==1,1]=100

Is it possible to do it more efficiently (w.r.t. lines of code). Also it would be nice to scale for additional dimensions in data_array (beyond 2d)

Upvotes: 1

Views: 93

Answers (1)

Alexander Korovin
Alexander Korovin

Reputation: 1475

If I understand your question correctly, please try this:

import numpy as np
index_array= np.array([0,1,1,0,0,0,1],dtype=bool)
data_array = np.zeros([len(index_array),2])
data_array[index_array,:]=100
print(data_array)

Here, index_array is instantiated as boolean. Or it can convert index_array to boolean using == (index_array==0).

Use data_array[index_array,...]=100 if there are additional dimensions.

------ CORRECTED VERSION ------

Hope I now understand your question. So try this:

import numpy as np
index_array= np.array([0,1,1,0,2,0,1,2])
data_array = np.zeros([len(index_array),3])

data_array[np.arange(len(index_array)), index_array] = 100

print(data_array)

Upvotes: 1

Related Questions