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

user_14492001 0
Viewed 45

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]])
pytorch indexing numpy arrays python
3 Answers

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]])
user_163434640

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]]
user_109729190

Since this question was tagged with pytorch, I will provide an alternative using torch.scatter_add. The effect is quite straightforward, calling:

x.scatter_add(dmi=1, index=idx, src=value)

will perform the following operation:

x[i][index[i]] += value[i][j]

So given x, idx, and c:

>>> x = torch.tensor([[0, 1, 0, 0, 0],
                      [0, 0, 0, 1, 0],
                      [1, 0, 0, 0, 0]])

>>> idx = torch.tensor([3, 1, 0])

Here is a one-line approach:

>>> x.scatter_add(1, idx[:,None], torch.full_like(x, c))
tensor([[0, 1, 0, 2, 0],
        [0, 2, 0, 1, 0],
        [3, 0, 0, 0, 0]])
user_63313690
Related