python - 将 pandas 条目连接到单个列列表中

Question

这是一个 pandas DataFrame 示例:

import pandas as pd
import numpy as np

dict_example = {'col1':['A', np.NaN, 'A', np.NaN, 'A'], 
    'col2':[np.NaN, 'B', 'A', np.NaN, 'A'], 'col3':['A', np.NaN, np.NaN, 'C', 'B'], 'col4':['D', np.NaN, 'B', 'D', 'A']}

df = pd.DataFrame(dict_example)
print(df)

  col1 col2 col3 col4
0    A  NaN    A    D
1  NaN    B  NaN  NaN
2    A    A  NaN    B
3  NaN  NaN    C    D
4    A    A    B    A

我想创建一个数据框列，其中包含行中没有 NaN 值的所有列的列表，即这个

print(df['new_column'])

['col1', 'col3', 'col4']
['col2']
['col1', 'col2', 'col4']
['col3', 'col4']
['col1', 'col2', 'col3', 'col4']
Name: new, dtype: object

我尝试使用以下函数和 apply() 来创建包含列表中所有元素的列。

def put_items_in_list(row):
    return [row['col1'], row['col2'], row['col3'], row['col4']]

df['new'] = df.apply(lambda row: put_items_in_list(row), axis=1)

print(df['new'])

0        [A, nan, A, D]
1    [nan, B, nan, nan]
2        [A, A, nan, B]
3      [nan, nan, C, D]
4          [A, A, B, A]
Name: new, dtype: object

但是，使用这个 lambda apply() 函数，我需要手动指定每个列名称。

此外，我想跟踪非 NaN 行值的列名称，而不是行值。为了做到这一点，需要有很多 if 语句，这感觉计算成本很高。

Answer 1

您可以stack和groupby:

df['new_col'] = df.stack().reset_index().groupby('level_0').level_1.agg(list)

或申请:

df['new_col'] = df.apply(lambda x: list(x.dropna().index), axis=1)

输出:

level_0
0          [col1, col3, col4]
1                      [col2]
2          [col1, col2, col4]
3                [col3, col4]
4    [col1, col2, col3, col4]
Name: level_1, dtype: object