CODEWITHSUNDEEP
pythonarraysnumpyindexing
honey
honey

Reputation: 13

How to index a numpy array with another numpy array in python

I am trying to index an np.array with another array so that I can have zeros everywhere after a certain index but it gives me the error

TypeError: only integer scalar arrays can be converted to a scalar index

Basically what I would like my code to do is that if I have:

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
d = np.array([2, 1, 3])

that I could do something like

a[d:] = 0

to give the output

a = [[ 1  2  3]
     [ 4  0  6]
     [ 0  0  9]
     [ 0  0  0]]

Upvotes: 1

Views: 1493

Answers (2)

hpaulj
hpaulj

Reputation: 231355

A widely used approach for this kind of problem is to construct a boolean mask, comparing the index array with the appropriate arange:

In [619]: mask = np.arange(4)[:,None]>=d
In [620]: mask
Out[620]: 
array([[False, False, False],
       [False,  True, False],
       [ True,  True, False],
       [ True,  True,  True]])
In [621]: a[mask]
Out[621]: array([ 5,  7,  8, 10, 11, 12])
In [622]: a[mask] = 0
In [623]: a
Out[623]: 
array([[1, 2, 3],
       [4, 0, 6],
       [0, 0, 9],
       [0, 0, 0]])

That's not necessarily faster than a row (or in this case column) iteration. Since slicing is basic indexing, it may be faster, even if done several times.

In [624]: for i,v in enumerate(d):
     ...:     print(a[v:,i])
     ...: 
[0 0]
[0 0 0]
[0]

Generally if a result involves multiple arrays or lists with different lengths, there isn't a "neat" multidimensional solution. Either iterate over those lists, or step back and "think outside the box".

Upvotes: 0

Tls Chris
Tls Chris

Reputation: 3824

It can be done with array indexing but it doesn't feel natural.

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
d = np.array([2, 1, 3])

col_ix = [ 0, 0, 1, 1, 1, 2 ]  # column ix for each item to change                                   
row_ix = [ 2, 3, 1, 2, 3, 3 ]  # row index for each item to change

a[ row_ix, col_ix ] = 0

a 
# array([[1, 2, 3],
#        [4, 0, 6],
#        [0, 0, 9],
#        [0, 0, 0]])

With a for loop

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])

for ix_col, ix_row in enumerate( d ):  # iterate across the columns
    a[ ix_row:, ix_col ] = 0

a 
# array([[1, 2, 3],
#        [4, 0, 6],
#        [0, 0, 9],
#        [0, 0, 0]])

Upvotes: 1

Related Questions