How to iterate over a sparse matrix with condition in Python

Question

I have a matrix with only 0 and 1 and it works as the indicator. Now I would like to replace the element of the matrix with value of 5 to the power of the original element. (Python please)

For example, if: old element = 0 -> new element is 5^0 = 1; else old element = 1 -> new element is 5^1 = 5

INPUT:

0 1 0
0 1 1

OUTPUT:

1 5 1
1 5 5

Please help this is urgent for my thesis code and this is also my first (real) time posting a question. I am trying to find power function or exponential package but it is not the one I need.

Thank you so much!

Answer 1

No need to iterate (in python code), just use your 0/1 array as index

 arr = np.array([[0,1,0],[1,0,1]])
 new = np.array([1,5])[arr]