Adding two numeric pandas columns with different lengths based on condition


I am writing a piece of simulation software in python using pandas, here is my problem:

Imagine you have two pandas dataframes dfA and dfB with numeric columns A and B respectively.
Both dataframes have a different number of rows denoted by n and m.
Let’s assume that n > m.
Moreover, dfA includes a binary column C, which has m times 1, and the rest 0.
Assume both dfA and dfB are sorted.

My question is, in order, I want to add the values in B to the values in column A if column C == 0.

In the example n = 6, m = 3.

Example data:

dataA = {'A': [7,7,7,7,7,7],
         'C': [1,0,1,0,0,1]}
dfA = pd.Dataframe(dataA)
dfB = pd.Dataframe([3,5,4], columns = ['B'])

Example pseudocode:

if dfA['C'] == 1:
    dfD['D'] = dfA['A']
    dfD['D'] = dfA['A'] + dfB['B']

Expected result:


I can only think of obscure for loops with index counters for each of the three vectors, but I am sure that there is a faster way by writing a function and using apply. But maybe there is something completely different that I am missing.

*NOTE: In the real problem the rows are not single values, but row vectors of equal length. Moreover, in the real problem it is not just simple addition but a weighted average over the two row vectors

Asked By: jorisvd



You can use:

m = dfA['C'].eq(1)
dfA['C'] = dfA['A'].where(m, dfA['A']+dfB['B'].set_axis(dfA.index[~m]))


dfA.loc[m, 'C'] = dfA.loc[m, 'A']
dfA.loc[~m, 'C'] = dfB['B'].values


   A   C
0  7   7
1  7  10
2  7   7
3  7  12
4  7  11
5  7   7
Answered By: mozway

The alternative answer is pretty clever. I am just showing a different way if you would like to do it using loops:

# Create an empty df
dfD = pd.DataFrame() 

# Create Loop
k = 0
for i in range(len(dfA)):
  if dfA.loc[i, "C"] == 1:
    dfD.loc[i, "D"] = dfA.loc[i, "A"]
   dfD.loc[i, "D"] = dfA.loc[i, "A"] + dfB.loc[k, "B"]
   k = k+1

# Show results

enter image description here

Answered By: Ledian K.
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.