Python中分组统计跨列唯一序列计数及缺失值处理问题

问题背景 #

有一个包含x、y1、y2、y3、y4列的Pandas DataFrame A，需要统计每个x对应的{y1,y2,y3,y4}唯一序列的出现次数。目前通过循环拼接子数据、去重得到了包含各x唯一序列的DataFrame B，但使用apply统计次数时，含NaN的序列count值为0——因为NaN与任何值（包括自身）比较都会返回False，无法满足将缺失值视为唯一值的需求。

现有问题代码 #

B = pd.DataFrame()
for x_temp in A['x'].unique():
    B = pd.concat([B, A[A['x'] == x_temp][['x','y1','y2','y3','y4']]])
B = B.drop_duplicates().sort_values(by=['x','y1','y2','y3','y4'])
del x_temp 

B['count'] = A.apply(lambda row: (A['y1'] == row['y1']) & (A['y2'] == row['y2']) & (A['y3'] == row['y3']) & (A['y4'] == row['y4']), axis=1).sum()

示例数据 #

import pandas as pd
import numpy as np

A = pd.DataFrame({'x':['1','1','1','1','2','2','2','2','2','1'],
                   'y1':['1','2','2',np.nan,'2',np.nan,'2','2','2','1'],
                   'y2':['2','1','2','2',np.nan,np.nan,'2','2','1','2'],
                   'y3':['1','1',np.nan,'2',np.nan,'2',np.nan,np.nan,'1','1'],
                   'y4':['2','2','2',np.nan,np.nan,'1','2','2','2','2']})

解决方案 #

方法一：直接分组统计（推荐） #

利用Pandas的groupby功能，设置dropna=False将NaN视为独立分组键，直接得到包含唯一序列和对应计数的结果：

# 按x和y1-y4分组，统计每个序列的出现次数
B = A.groupby(['x', 'y1', 'y2', 'y3', 'y4'], dropna=False).size().reset_index(name='count')
# 按需求排序
B = B.sort_values(by=['x','y1','y2','y3','y4'])

方法二：基于原有B逻辑的改进版 #

如果需要保留原有生成B的逻辑，可以先得到唯一序列的B，再通过合并分组统计结果添加count列：

# 生成唯一序列的B（简化原有循环逻辑）
B = A.drop_duplicates(subset=['x','y1','y2','y3','y4']).sort_values(by=['x','y1','y2','y3','y4'])
# 统计每个序列的次数
count_df = A.groupby(['x', 'y1', 'y2', 'y3', 'y4'], dropna=False).size().reset_index(name='count')
# 合并count到B
B = B.merge(count_df, on=['x','y1','y2','y3','y4'], how='left')

原理说明 #

原有代码失效的核心原因是NaN的比较特性：np.nan == np.nan会返回False，导致apply中的条件判断无法匹配含NaN的序列。而groupby的dropna=False参数会将NaN当作独特的分组键，正确统计其出现次数，完美匹配需求。

内容的提问来源于stack exchange，提问作者Joshua Roy