Add a scalar to a numpy matrix based on the indices in a different numpy array

user_14762717 0
Viewed 405

I'm sorry if this question isn't framed well. So I would rather explain with an example. I have a numpy matrix:

a = np.array([[0.5, 0.8, 0.1], [0.6, 0.9, 0.3], [0.7, 0.4, 0.8], [0.8, 0.7, 0.6]])

And another numpy array as shown:

b = np.array([1, 0, 2, 2])

With the given condition that values in b will be in the range(a.shape[1]) and that b.shape[1] == a.shape[0]. Now this is the operation I need to perform.

For every index i of a, and the corresponding index i of b, I need to subtract 1 from the index j of a[i] where j== b[i]

So in my example, a[0] == [0.5 0.8 0.1] and b[0] == 1. Therefore I need to subtract 1 from a[0][b[0]] so that a[0] = [0.5, -0.2, 0.1]. This has to be done for all rows of a. Any direct solution without me having to iterate through all rows or columns one by one?

Thanks.

numpy arrays python
1 Answers

Use numpy indexing. See this post for a nice introduction:

import numpy as np

a = np.array([[0.5, 0.8, 0.1], [0.6, 0.9, 0.3], [0.7, 0.4, 0.8], [0.8, 0.7, 0.6]])
b = np.array([1, 0, 2, 2])

a[np.arange(a.shape[0]), b] -= 1

print(a)

Output

[[ 0.5 -0.2  0.1]
 [-0.4  0.9  0.3]
 [ 0.7  0.4 -0.2]
 [ 0.8  0.7 -0.4]]

As an alternative use substract.at:

np.subtract.at(a, (np.arange(a.shape[0]), b), 1)
print(a)

Output

[[ 0.5 -0.2  0.1]
 [-0.4  0.9  0.3]
 [ 0.7  0.4 -0.2]
 [ 0.8  0.7 -0.4]]

The main idea is that:

np.arange(a.shape[0])  # shape[0] is equals to the numbers of rows

generates the indices of the rows:

[0 1 2 3]
user_40015920
Related