如何找到列号，并返回它作为一个数组时，在python列中有一个值?

Question

假设您有以下数据帧:

ID    Q1       Q2      Q3
0     1        0        1
1     0        0        1
2     1        1        1
3     0        1        1  

我想返回一个列编号的数组，只要有一个1，并添加它作为另一列，像这样:

ID   Array    Q1       Q2      Q3
0    [0,2]     1        0       1
1    [2]       0        0       1
2    [0,1,2]   1        1       1
3    [1,2]     0        1       1 

感谢

Answer 1

我将使用numpy.where:

a, b = np.where(df.filter(like='ID')==1)
# or
a, b = np.where(df.drop(columns='ID')==1)
df['Array'] = pd.Series(b).groupby(a).agg(list).set_axis(df.index)

输出:

ID  Q1  Q2  Q3      Array
0   0   1   0   1     [0, 2]
1   1   0   0   1        [2]
2   2   1   1   1  [0, 1, 2]
3   3   0   1   1     [1, 2]

纯熊猫变体:

df2 = df.filter(like='Q')
df['Array'] = (df2.set_axis(range(df2.shape[1]), axis=1).stack()
.loc[lambda s: s==1].reset_index()
.groupby('level_0')['level_1'].agg(list)
)

Answer 2

另一种使用np的方法。凡与"应用"在一起的都是"应用"。方法:

import pandas as pd
import numpy as np
data = {'ID': {0: 0, 1: 1, 2: 2, 3: 3},
'Q1': {0: 1, 1: 0, 2: 1, 3: 0},
'Q2': {0: 0, 1: 0, 2: 1, 3: 1},
'Q3': {0: 1, 1: 1, 2: 1, 3: 1}}
df = pd.DataFrame(data)
#####
def arr_func(row):
return np.where(row)[0]
df['Array'] = df.drop(columns = 'ID').apply(arr_func, axis = 1)

结果:

ID  Q1  Q2  Q3      Array
0   0   1   0   1     [0, 2]
1   1   0   0   1        [2]
2   2   1   1   1  [0, 1, 2]
3   3   0   1   1     [1, 2]

Answer 3

这里有一个简单的解决方案。

import pandas as pd
def my_func(row):
return [item for item in row if item == 1] 
df['Array'] = df[['Q1', 'Q2', 'Q3']].apply(my_func, axis=1)