CODEWITHSUNDEEP
numpy
zdm
zdm

Reputation: 531

One line to replace values with their row indices?

I have a numpy array

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

I would like to replace every positive elements of this array by its row index+1. So the final result would be:

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

Can I do this with a simply numpy command (without looping)?

Upvotes: 0

Views: 98

Answers (2)

Chris
Chris

Reputation: 29742

Use numpy.arange

(a != 0) * np.reshape(np.arange(a.shape[0])+1, (-1, 1))

Output:

array([[1., 0., 0., 1., 0.],
       [0., 2., 0., 0., 0.],
       [0., 0., 3., 0., 3.]])

Works on any array:

a2 = np.array([[1,0,0,-1,0],
               [0,20,0,0,0],
               [0,0,-300,0,30]])
(a2 != 0) * np.reshape(np.arange(a2.shape[0])+1, (-1, 1))

Output:

array([[1., 0., 0., 1., 0.],
       [0., 2., 0., 0., 0.],
       [0., 0., 3., 0., 3.]])

Upvotes: 2

tobias_k
tobias_k

Reputation: 82899

Not sure if this is the proper numpy way, but you could use enumerate and multiply the sub-arrays by their indices:

>>> np.array([x * i for i, x in enumerate(a, start=1)])
array([[1, 0, 0, 1, 0],
       [0, 2, 0, 0, 0],
       [0, 0, 3, 0, 3]])

Note that this only works properly if "every positive element" is actually 1, as in your example, and if there are no negative elements. Alternatively, you can use a > 0 to first get an array with True (i.e. 1) in every place where a is > 0 and False (i.e. 0) otherwise.

>>> a = np.array([[ 1, 0, 0, 2, 0],
...               [ 0, 3, 0, 0,-8],
,,,               [-3, 0, 4, 0, 5]])
...
>>> np.array([x * i for i, x in enumerate(a > 0, start=1)])
array([[1, 0, 0, 1, 0],
       [0, 2, 0, 0, 0],
       [0, 0, 3, 0, 3]])

Upvotes: 1

Related Questions