按分组获取Pandas DataFrame最小vd的NaN与非NaN行

原始DataFrame #

import pandas as pd
import numpy as np

df = pd.DataFrame({'id' : [1,1,1,1,1,2,2,2,2,2,3,3,3,3,3],
                   'vd' : [1, 2, 3, 4, 5, 2, 3, 4, 5, 6,
                           3, 4, 5, 6, 7],
            'value'  : [np.NaN, np.NaN, np.NaN, 2.3, 1.1,
                        np.NaN, np.NaN, 2.2, 3.3, 2.2,
                        np.NaN, 1, 2, 2.3, 1.1]})

需求说明 #

按id分组后，为每个id筛选两类行：

• vd最小且value为NaN的行
• vd最小且value不为NaN的行

期望输出：

df_out = pd.DataFrame({'id' : [1,1, 2,2, 3,3],
                   'vd' : [1, 4, 2, 4,
                           3, 4],
            'value'  : [np.NaN, 2.3,
                        np.NaN, 2.2,
                        np.NaN, 1]})

已尝试方法及问题 #

已尝试用以下代码筛选非NaN的最小vd行：

out = df[df.index == (df['value'].notnull().groupby(df['id']).transform('idxmax'))]

但在获取NaN行的最小vd时，使用idxmin无法得到正确索引：

df['value'].isnull().groupby(df['id']).transform('idxmin')

问题在于isnull()返回布尔值，idxmin会返回每组中第一个False的索引，而非我们需要的NaN行里vd最小的行索引。

解决方案 #

直接拆分两类行分别处理后合并，逻辑更清晰：

# 筛选value为NaN的行，按id分组取vd最小的行
subset_nan = df[df['value'].isnull()].groupby('id').apply(lambda x: x[x['vd'] == x['vd'].min()]).reset_index(drop=True)

# 筛选value不为NaN的行，按id分组取vd最小的行
subset_not_nan = df[df['value'].notnull()].groupby('id').apply(lambda x: x[x['vd'] == x['vd'].min()]).reset_index(drop=True)

# 合并两个子集并按id排序
df_out = pd.concat([subset_nan, subset_not_nan]).sort_values('id').reset_index(drop=True)

或者用更高效的idxmin直接获取每组vd最小的索引，再分别筛选：

# 处理NaN行：先过滤出NaN行，按id取vd最小的索引
nan_idx = df[df['value'].isnull()].groupby('id')['vd'].idxmin()
subset_nan = df.loc[nan_idx]

# 处理非NaN行：过滤出非NaN行，按id取vd最小的索引
not_nan_idx = df[df['value'].notnull()].groupby('id')['vd'].idxmin()
subset_not_nan = df.loc[not_nan_idx]

# 合并排序
df_out = pd.concat([subset_nan, subset_not_nan]).sort_values('id').reset_index(drop=True)

这两种方法都能得到符合预期的输出。

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