Pandas：对于每一行，获取匹配的列

Question

我有一个这样的数据框，

df = pd.DataFrame({'A':[1,-2,3,4,5],'B':[1,2,3,4,-5],'C':[1,2,3,-4,5],'value':[123,1567,1456,764,2456]})
print(df)

是，

   A  B  C  value
0  1  1  1    123
1 -2  2  2   1567
2  3 3  3   1456
3  4  4 -4    764
4  5  -5  5   2456

我知道每一行只有一个负值。我想知道这样的值在哪一列。我知道我可以这样做，

desiredOutput = []
for i,row in df.iterrows():
    if any(row < 0):
        print(row[row < 0].index.to_numpy()[0],row[row[row < 0].index.to_numpy()[0]],row.value)
        desiredOutput.append(row[row < 0].index.to_numpy()[0])
    else:
        desiredOutput.append(None)

print(desiredOutput)

给予，

A -2 1567
C -4 764
B -5 2456
[None, 'A', None, 'C', 'B']

但我想一定有一种 pythonnic 的方式来做到这一点（可能使用 .apply()？）

Answer 1

使用DataFrame.dot进行矩阵乘法：

desiredOutput = df.lt(0).dot(df.columns).mask(lambda x: x.eq(''), None)
print (desiredOutput)
0    None
1       A
2    None
3       C
4       B
dtype: object

desiredOutput = df.lt(0).dot(df.columns).replace('',np.nan)
print (desiredOutput)
0    NaN
1      A
2    NaN
3      C
4      B
dtype: object

使用 apply 是可能的，但如果按每行的掩码过滤列会更慢：

desiredOutput = df.lt(0).apply(lambda x: next(iter(x.index[x]), None), axis=1)
print (desiredOutput)
0    None
1       A
2    None
3       C
4       B
dtype: object

Answer 2

您可以使用 idxmax on a boolean mask and where 来屏蔽没有值匹配的情况：

cols = ['A', 'B', 'C']  # optional: if needed to select columns
m = df[cols].lt(0)      # else:  df.lt(0)
df['result'] = m.idxmax(axis=1).where(m.any(1))

输出：

   A  B  C  value result
0  1  1  1    123    NaN
1 -2  2  2   1567      A
2  3  3  3   1456    NaN
3  4  4 -4    764      C
4  5 -5  5   2456      B

替代所有数值列的 one-liner（要求 python ≥3.8）：

df['result'] = ((m:=df.select_dtypes('number').lt(0))
                .idxmax(axis=1).where(m.any(1))
                )

Pandas：对于每一行，获取匹配的列

Pandas: for each row, get the matching column

python

pandas

apply