根据pandas中的第三列保留两列之间的值

Question

I have three columns, A, B and C. I want to create a fourth column D that contains values of A or B, based on the value of C. For example:

   A   B   C   D 
0  1   2   1   1
1  2   3   0   3
2  3   4   0   4
3  4   5   1   4

In the above example, column D takes the value of column A if the value of C is 1 and the value of column B if the value of C is 0. Is there an elegant way to do it in Pandas? Thank you for your help.

Answer 1

Use numpy.where :

In [20]: df
Out[20]:
   A  B  C
0  1  2  1
1  2  3  0
2  3  4  0
3  4  5  1

In [21]: df['D'] = np.where(df.C, df.A, df.B)

In [22]: df
Out[22]:
   A  B  C  D
0  1  2  1  1
1  2  3  0  3
2  3  4  0  4
3  4  5  1  4

Answer 2

pandas
In consideration of the OP's request

Is there an elegant way to do it in Pandas?

my opinion of elegance
and idiomatic pure pandas
assign + pd.Series.where

df.assign(D=df.A.where(df.C, df.B))

   A  B  C  D
0  1  2  1  1
1  2  3  0  3
2  3  4  0  4
3  4  5  1  4

---

response to comment

how would you modify the pandas answer if instead of 0, 1 in column C you had A, B?

df.assign(D=df.lookup(df.index, df.C))

   A  B  C  D
0  1  2  A  1
1  2  3  B  3
2  3  4  B  4
3  4  5  A  4