I have a dataframe below
df=pd.DataFrame(np.random.randn(6,3),index=list("ABCDEF"),columns=list("XYZ"))
df.reset_index(inplace=True)
df
I want to have a new column named "Q". The values under column "Q" shall be calculated based on the labels under index column with the following three conditions:
conditions=[(df["index"]== "A"|"B"|"C"|"D"),(df["index"]== "E"),(df["index"]== "F")]
returned_value=[df["X"]+df["Y"],df["Y"]*2,df["Z"]]
So I was thinking using
df["Q"]=np.select(conditions, returned_value)
I however got the error after defining the conditions. I first used or, and got another error, and then changed to |, but got the following. Any hints on how can I achieve what I want?
TypeError: unsupported operand type(s) for |: 'str' and 'str'
Use isin for check membership of multiple values:
np.random.seed(1213)
df=pd.DataFrame(np.random.randn(6,3),index=list("ABCDEF"),columns=list("XYZ"))
df.reset_index(inplace=True)
conditions=[df["index"].isin(["A","B","C","D"]),(df["index"]== "E"),(df["index"]== "F")]
returned_value=[df["X"]+df["Y"],df["Y"]*2,df["Z"]]
df["Q"]=np.select(conditions, returned_value)
print (df)
index X Y Z Q
0 A 0.511604 -0.217660 -0.521060 0.293943
1 B 1.253270 1.104554 -0.770309 2.357825
2 C 0.632975 -1.322322 -0.936332 -0.689347
3 D 0.436361 1.233744 0.527565 1.670105
4 E -0.369576 1.820059 -1.373630 3.640118
5 F -0.414554 -0.098443 0.904791 0.904791
But reset index is not necessary, then check df.index:
np.random.seed(1213)
df=pd.DataFrame(np.random.randn(6,3),index=list("ABCDEF"),columns=list("XYZ"))
conditions=[df.index.isin(["A","B","C","D"]),(df.index == "E"),(df.index== "F")]
returned_value=[df["X"]+df["Y"],df["Y"]*2,df["Z"]]
df["Q"]=np.select(conditions, returned_value)
print (df)
X Y Z Q
A 0.511604 -0.217660 -0.521060 0.293943
B 1.253270 1.104554 -0.770309 2.357825
C 0.632975 -1.322322 -0.936332 -0.689347
D 0.436361 1.233744 0.527565 1.670105
E -0.369576 1.820059 -1.373630 3.640118
F -0.414554 -0.098443 0.904791 0.904791