CODEWITHSUNDEEP
pythonnumpy
Philipp
Philipp

Reputation: 62

Apply a function based on condition and column index position

I have a 2d numpy array for which I would like to apply a computation on elements based on a condition and considering the column index position.

The condition is: element value == 0.

If true I want to compute 10^i (where i is the index of axis 1) in that position otherwise 0.

So an example array could be:

ar = np.array([[  0,  -5,  -5, -11,  -9],
               [  5,   0,   0,  -6,  -4],
               [ 10,   5,   5,  -1,   1]])

And the resulting array should look like this (with the exponents evaluated):

[[  10^0,  0,  0, 0,  0],
 [  0,   10^1,   10^2,  0,  0],
 [ 0,   0,   0,  0,   0]]

I've gotten as far as:

np.where(ar==0)

I'm not sure how to apply the column index as exponent given that the condition is true.

Thank you!

Upvotes: 2

Views: 118

Answers (1)

amzon-ex
amzon-ex

Reputation: 1744

Create test array:

a = np.random.choice(range(-4, 4), size = (3, 4))

# a = array([[ 0, -1, -1, -2],
#            [ 0, -3,  0, -1],
#            [ 0,  3, -2,  0]])

Find indices where a is zero:

mask = a == 0
indx = np.nonzero(mask)

# indx = (array([0, 1, 1, 2, 2]), array([0, 0, 2, 0, 3]))

Change elements:

a[mask] = 10**(indx[1])
a[~mask] = 0

# a = array([[   1,    0,    0,    0],
#            [   1,    0,  100,    0],
#            [   1,    0,    0, 1000]])

indx is a tuple whose second element stores column indices for a 2-D array.

Upvotes: 3

Related Questions