检查行中是否存在值并提取列的名称- Pandas

Question

我有一个数据框架如下:

df = 
A          col_1     col_45    col_9    col_10

1.0     4.0        45.0      NaN       34.9     NaN
2.0     4.0        NaN       NaN       23.4     45.6      
3.0     49.2       10.8      NaN       NaN      23.8

对于数据框中的每一行，我想检查col_1、col_45、col_9和col_10是否存在任何值。如果存在值，我想获得列名的数字，并将其放入列表/数组中。例如,

在第一行中，有col_1和col_9的值。所以我想从列名中获得数字1和9，并放入列表/数组。这应该为每一行执行。

是否有简单的方法来做到这一点?

Answer 1

使用DataFrame.drop和DataFrame.melt进行unpivot，删除缺失的值，将列名转换为整数，并聚合lists:

df['new'] = (df.drop('A',1)
.melt(ignore_index=False, var_name='a')
.dropna().a
.str.replace('D','', regex=True)
.astype(int)
.groupby(level=0)
.agg(list))
print (df)
A  col_1  col_45  col_9  col_10      new
1.0   4.0   45.0     NaN   34.9     NaN   [1, 9]
2.0   4.0    NaN     NaN   23.4    45.6  [9, 10]
3.0  49.2   10.8     NaN    NaN    23.8  [1, 10]

Answer 2

对于列名，您可以使用stack(自动删除所有NaN)，然后使用GroupBy.agg来形成列表:

out = (df.set_index('A', append=True)
.stack().reset_index(-1)
.groupby(level=0)['level_2'].agg(list)
)

输出:

1.0     [col_1, col_9]
2.0    [col_9, col_10]
3.0    [col_1, col_10]
Name: level_2, dtype: object

数字:

out = (df.set_index('A', append=True)
.stack().reset_index(-1)
['level_2'].str.extract('_(d+)', expand=False).astype(int)
.groupby(level=0).agg(list)
)

输出:

1.0     [1, 9]
2.0    [9, 10]
3.0    [1, 10]
Name: level_2, dtype: object

另一种方法，你可以使用掩码和矩阵乘法:

cols = df.filter(like='col').columns
# Index(['col_1', 'col_45', 'col_9', 'col_10'], dtype='object')
int_cols = cols.str.extract('_(d+)$', expand=False).astype(int)
# Int64Index([1, 45, 9, 10], dtype='int64')
m = df[cols].notna()
#      col_1  col_45  col_9  col_10
# 1.0   True   False   True   False
# 2.0  False   False   True    True
# 3.0   True   False  False    True
out = (m.astype(int).mul(int_cols).where(m).convert_dtypes()
.stack().groupby(level=0).agg(list)
)
# 1.0     [1, 9]
# 2.0    [9, 10]
# 3.0    [1, 10]
# dtype: object