2

I have data stored in three different arrays: A, B and C.
I also have an integer array, Z, that has either 1, 2 or 3 in it.
A, B, C and Z have the same shape.
I want to create a new array, D, that contains the value listed in A if the corresponding element in Z is 1, the value listed in B if the corresponding element in Z is 2 or the value listed in C if the corresponding element in Z is 3.

To do this, I have written code with nested versions of numpy.where. The code however looks ugly. Is there a better way to do the same ?

import numpy as np
#Create arrays that hold the data
A = np.array([1.,1.,1.])
B = A * 2
C = A * 3
#Create an array that hold the zone numbers
Z = np.array([1,2,3])
#Now calculate the new array by figuring out the appropriate zone at each location
D = np.where(Z==1,A,np.where(Z==2,B, np.where(Z==3,C,0.0)))
#Output
print A
print B
print C
print D

[ 1.  1.  1.]
[ 2.  2.  2.]
[ 3.  3.  3.]
[ 1.  2.  3.]
Prashanth
  • 115
  • 2
  • 10

1 Answers1

2

See if this works:

new=np.zeros_like(A)
new[Z==1] = A[Z == 1]
new[Z==2] = B[Z == 2]
new[Z==3] = C[Z == 3]
wwii
  • 23,232
  • 7
  • 37
  • 77
  • Thanks. I think I will adopt this approach to improve code readability. – Prashanth Jul 01 '14 at 20:22
  • check the `np.select` solution, which is way faster. https://stackoverflow.com/questions/49252969/alternative-to-nested-np-where-in-pandas-dataframe. – Jia Gao Sep 20 '21 at 01:22
  • @JasonGoal - you should write an answer using `.select` and maybe add some example data that can be used for performance testing - maybe the OP will change their mind and select your answer.. – wwii Sep 20 '21 at 13:58