Numpy, increment values in a 2D array using index represented in another 1D array-CodePudding

Here is an example of what I would like to do: Assume Array A

A = np.array([[0, 1, 3, 5, 9], 
              [2, 7, 5, 1, 4]])

And Array B

B = np.array([2, 4])

I am looking for an operation that will increment the element indexed by array B in each row of array A by 1. So the result A is:

A = np.array([[0, 1, 4, 5, 9], 
              [2, 7, 5, 1, 5]])

The index 2 of first row is increased by 1, and the index 4 of second row is increased by 1

CodePudding user response：

You can achieve this by using advanced indexing in numpy:

A[np.arange(len(B)), B]  = 1

This works by creating a 2D array with dimensions (len(B), len(B)) using np.arange(len(B)), which represents the row indices. The second index of the advanced indexing, B, represents the column indices. By adding 1 to A[np.arange(len(B)), B], you increment the elements in each row specified by B.

CodePudding user response：

In numpy you can do by using arrange and shape of an array

import numpy as np

A = np.array([[0, 1, 3, 5, 9], 
              [2, 7, 5, 1, 4]])
B = np.array([2, 4])

A[np.arange(A.shape[0]), B]  = 1

print(A)

np.arange(A.shape[0]) generates an array of integers from 0 to A.shape[0] - 1. A.shape[0] is basically rows

you can do with looping also..

import numpy as np
A = np.array([[0, 1, 3, 5, 9], 
              [2, 7, 5, 1, 4]])
B = np.array([2, 4])

for i, index in enumerate(B):
    A[i][index]  = 1

print(A)