如何按特定列分组，然后在不是NA的地方获取多个列的计数，并将它们添加到Pandas Python中？

import pandas as pd

# 创建一个示例DataFrame
data = {'A': [1, 2, 3, 4, 5],
        'B': [1, 1, 2, 2, 2],
        'C': [1, 2, pd.NA, 4, pd.NA],
        'D': [1, pd.NA, pd.NA, 4, 5]}
df = pd.DataFrame(data)

# 按列'B'进行分组，并计算列'C'和列'D'的非NA值的计数
result = df.groupby('B').agg({'C': lambda x: x.count(), 'D': lambda x: x.count()})

# 将计数结果添加到原DataFrame中
df['C_count'] = df['B'].map(result['C'])
df['D_count'] = df['B'].map(result['D'])

print(df)

   A  B    C    D  C_count  D_count
0  1  1    1    1        2        2
1  2  1    2  NaN        2        2
2  3  2  NaN  NaN        3        3
3  4  2    4    4        3        3
4  5  2  NaN    5        3        3