python - 在python数据帧中的每列的最大值之前找到一个值的索引

Question

我有一个数据框如下。

test = pd.DataFrame({'col1':[0,0,1,0,0,0,1,2,0], 'col2': [0,0,1,2,3,0,0,0,0]})
   col1  col2
0     0     0
1     0     0
2     1     1
3     0     2
4     0     3
5     0     0
6     1     0
7     2     0
8     0     0

对于每一列，我想在每一列的最大值之前找到值为 1 的索引。例如，对于第一列，最大值为2，值1在2之前的索引为6。对于第二列，最大值为3，值1在值3之前的索引为2。

总而言之，我希望得到 [6, 2] 作为这个测试 DataFrame 的输出。有没有快速的方法来实现这一目标？

Answer 1

使用Series.mask隐藏不是 1 的元素，然后应用 Series.last_valid_index到每一列。

m = test.eq(test.max()).cumsum().gt(0) | test.ne(1) 
test.mask(m).apply(pd.Series.last_valid_index)

col1    6
col2    2
dtype: int64

---

用numpy向量化，可以用numpy.cumsum和 argmax :

idx = ((test.eq(1) & test.eq(test.max()).cumsum().eq(0))
            .values
            .cumsum(axis=0)
            .argmax(axis=0))
idx
# array([6, 2])

pd.Series(idx, index=[*test])

col1    6
col2    2
dtype: int64