Home > OS >  How to index an array/tensor on the "rest" of the indices given some indices?
How to index an array/tensor on the "rest" of the indices given some indices?

Time:07-29

Given some array (or tensor):

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

and some indices of dimension equaling the number of rows in x:

idx = np.array([3, 1, 0])  # Index values range from (0: number of columns) in "x"!

Now if I wanted to add a certain value c to the columns of x depending on the indices idx, I would do the following:

x[range(3), idx]  = c

To get:

x = np.array([[0, 1, 0, c, 0],
              [0, c, 0, 1, 0],
              [1 c, 0, 0, 0, 0]])

But what if I wanted to add the value c to every other column index in x, rather than the exact indices in idx?

The desired outcome (based on the above example) should be:

x = np.array([[c, 1 c, c, 0, c],
              [c, 0, c, 1 c, c],
              [1, c, c, c, c]])

CodePudding user response:

Create a boolean array to use as mask:

# set up default mask
m = np.ones(x.shape, dtype=bool)
# update mask
m[np.arange(m.shape[0]), idx] = False
# perform boolean indexing
x[m]  = c

Output (c=9):

array([[ 9, 10,  9,  0,  9],
       [ 9,  0,  9, 10,  9],
       [ 1,  9,  9,  9,  9]])

m:

array([[ True,  True,  True, False,  True],
       [ True, False,  True,  True,  True],
       [False,  True,  True,  True,  True]])

CodePudding user response:

Maybe you can use Masking matrix to multiply by value you need only section that does not intersect indices. Try this:

import numpy as np 

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

# initialize `c`
c = 9

idx = np.array([3, 1, 0])
#initialize mask 
mask = np.ones_like(x)
mask[range(3), idx] = 0 
mask *= c

x = x * mask

# output 

[[ 9 10  9  0  9]
 [ 9  0  9 10  9]
 [ 1  9  9  9  9]]
  • Related