Pandas - Sum of Column A where Column B is in Column C

Viewed 1198

I have the following dataframe. Notice that Column B is a series of lists. This is what is giving me trouble

Dataframe 1:

   Column A        Column B
0      10            [X]
1      20          [X,Y]
2      15        [X,Y,Z]
3      25            [A]
4      60            [B]

I want to take all of the values in Column C (below), check if they exist in Column B, and then sum their values from Column A.

DataFrame 2: (Desired Output)

   Column C        Sum of Column A
0       X                 45   
1       Y                 35
2       Z                 15
3       Q                  0
4       R                  0

I know this can be accomplished using a for-loop, but I am looking for the "pandonic method" to solve this.

3 Answers

update

Here is a shorter and faster answer beginning with your second dataframe

df2['C'].apply(lambda x: df.loc[df['B'].apply(lambda y: x in y), 'A'].sum())

original answer

You first can 'normalize' the data in Column B.

df_normal = pd.concat([df.A, df.B.apply(pd.Series)], axis=1)

    A  0    1    2
0  10  X  NaN  NaN
1  20  X    Y  NaN
2  15  X    Y    Z
3  25  A  NaN  NaN
4  60  B  NaN  NaN

And then stack and groupby to get the lookup table.

df_lookup = df_normal.set_index('A') \
                     .stack() \
                     .reset_index(name='group')\
                     .groupby('group')['A'].sum()

group
A    25
B    60
X    45
Y    35
Z    15
Name: A, dtype: int64

And then join to df2.

df2.join(df_lookup, on='C').fillna(0)

   C     A
0  X  45.0
1  Y  35.0
2  Z  15.0
3  Q   0.0
4  R   0.0

And in one line

df2.join(
    df.set_index('A')['B'] \
      .apply(pd.Series) \
      .stack() \
      .reset_index('A', name='group') \
      .groupby('group')['A'] \
      .sum(), on='C') \
   .fillna(0)

And if you wanted to loop which isn't that bad in this situation

d = {}
for _, row in df.iterrows():
    for var in row['B']:
        if var in d:
            d[var] += row['A']
        else:
            d[var] = row['A']

df2.join(pd.Series(d, name='Sum of A'), on='C').fillna(0)

Base on your example data:

df1=df.set_index('Column A')['Column B'].\
        apply(pd.Series).stack().reset_index().\
             groupby([0])['Column A'].sum().to_frame()
df2['Sum of Column A']=df2['Column C'].map(df1['Column A'])

df2.fillna(0)

Out[604]: 
  Column C  Sum of Column A
0        X             45.0
1        Y             35.0
2        Z             15.0
3        Q              0.0
4        R              0.0

Data input :

df = pd.DataFrame({'Column A':[10,20,15,25,60],'Column B':[['X'],['X','Y'],['X','Y','Z'],['A'],['B']]})
df2 = pd.DataFrame({'Column C':['X','Y','Z','Q','R']})

I'd use a list comprehension like

df['result']=np.sum[(df['Column C'] in col['Column B'])*col['Column A'] for col in df]
Related