根据计算值从Pandas DataFrame中的行中查找数据？

Question

As an extension of my previous question , I would like take a DataFrame like the one below and find the correct row from which to pull data from column C and place it into column D based upon the following criteria:

1. B_new = 2*A_old -B_old , ie. the new row needs to have a B equal to the following result from the old row: 2*A - B .
2. Where A is the same, ie. A in the new row should have the same value as the old row.
3. Any values not found should use a NaN result

Code:

import pandas as pd
a = [2,2,2,3,3,3,3]
b = [1,2,3,1,3,4,5]
c = [0,1,2,3,4,5,6]

df = pd.DataFrame({'A': a , 'B': b, 'C':c})
print(df)

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

Desired output:

   A  B  C    D
0  2  1  0  2.0
1  2  2  1  1.0
2  2  3  2  0.0
3  3  1  3  6.0
4  3  3  4  4.0
5  3  4  5  NaN
6  3  5  6  3.0

Based upon the solutions in my previous question , I've come up with a method that uses a for loop to move thru each unique value of A :

for i in df.A.unique():
    mapping = dict(df[df.A==i][['B', 'C']].values)
    df.loc[df.A==i,'D'] = (2 * df[df.A==i]['A'] - df[df.A==i]['B']).map(mapping)

However, this seem clunky and I suspect there is a better way that doesn't make use of for loops, which from my prior experience tend to be slow.

Question: What's the fastest way to accomplish this transfer of data within the DataFrame?

Answer 1

You could

In [370]: (df[['A', 'C']].assign(B=2*df.A - df.B)
           .merge(df, how='left', on=['A', 'B'])
           .assign(B=df.B)
           .rename(columns={'C_x': 'C', 'C_y': 'D'}) )
Out[370]:
   A  C  B    D
0  2  0  1  2.0
1  2  1  2  1.0
2  2  2  3  0.0
3  3  3  1  6.0
4  3  4  3  4.0
5  3  5  4  NaN
6  3  6  5  3.0

Details:

In [372]: df[['A', 'C']].assign(B=2*df.A - df.B)
Out[372]:
   A  C  B
0  2  0  3
1  2  1  2
2  2  2  1
3  3  3  5
4  3  4  3
5  3  5  2
6  3  6  1

In [373]: df[['A', 'C']].assign(B=2*df.A - df.B).merge(df, how='left', on=['A', 'B'])
Out[373]:
   A  C_x  B  C_y
0  2    0  3  2.0
1  2    1  2  1.0
2  2    2  1  0.0
3  3    3  5  6.0
4  3    4  3  4.0
5  3    5  2  NaN
6  3    6  1  3.0