Counting values with condition in one DataFrame and adding the result to another DataFrame

Viewed 58

I have two DataFrames:

df1 = pd.DataFrame({"id": [1, 2, 3, 4]})

df2 = pd.DataFrame({"id": [1, 1, 2, 4, 4, 4],
                  "text": ["a", "a", "b", "a", "b", "b"]})

Output df1:

   id
0   1
1   2
2   3
3   4

Output df2:

   id value
0   1    a
1   1    a
2   2    b
3   4    a
4   4    b
5   4    b

My goal is to add three columns in df1.

In count_all I would like to count the corresponding ids in df2. E.g. id 4 exists 3 times in df2.

In count_a I would like to count the corresponding ids in df2 where the text value == a.

In count_b I would like to count the corresponding ids in df2 where the text value == b.

   id  count_all  count_a  count_b
0   1          2        2        0
1   2          1        0        1
2   3          0        0        0
3   4          3        1        2

How can this be done with pandas?

2 Answers

Use crosstab with margins parameter, add missing index values or change columns ordering by DataFrame.reindex, change columns names by DataFrame.add_prefix and last join to df1 by DataFrame.join:

df = (df1.join(pd.crosstab(df2['id'], df2['text'], margins=True)
                 .reindex(index=df1['id'].unique(), 
                          columns=['All'] + df2['text'].unique().tolist(), 
                          fill_value=0)
                 .add_prefix('count_'), on='id'))
print (df)
   id  count_All  count_a  count_b
0   1          2        2        0
1   2          1        0        1
2   3          0        0        0
3   4          3        1        2

Here is another way:

df1.join(df2.groupby('id').agg(
    count_all = ('id','count'),
    count_a=('text',lambda x: sum(x.eq('a'))),
    count_b = ('text',lambda x: sum(x.eq('b')))),on='id').fillna(0)
Related