Indexing a 2D narray when row mask is a 1D list and column mask is a 2D list

user_19816922 0
Viewed 33

I have a 2D numpy array of zeros that looks like this:

stock2concept_matrix = np.zeros((3, 4))

I have a row mask, which is a list of indices from 0 to 2:

row = [0, 1, 2]

A column mask looks like this:

column = [[0, 2], [1], [0, 2, 3]]

I expect that

stock2concept_matrix[row, column] = 1

will change stock2concept_matrix to:

[[1, 0, 1, 0],
 [0, 1, 0, 0],
 [1, 0, 1, 1]]

However, an error occurred:

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

I don't want to use a for loop since it's inefficient, how can I achieve my goal?

numpy python
1 Answers

Use np.repeat on row to make the right number of elements for each column entry:

row = np.repeat(row, [len(c) for c in column])

Then flatten the column list to get a pair of indices like the result of np.where:

column = [x for c in column for x in c]

Now you can just do

stock2concept_matrix[row, column] = 1

Numpy indices must be broadcastable together. In this case, that means that you have to have one element in each index array per element that you want to set.

user_29887300
Related