Python 如果条件不满足，我如何获得掩码的第一个索引并获得None

发布于04月09日

这是我的DataFrame:

import pandas as pd
import numpy as np
df = pd.DataFrame(
    {
        'a': [np.nan, np.nan, 1, 2, 3, 1, 0, 1, -1, -1],
    }
)

我的面具是:

mask = (df.a.gt(2))

而expectet的输出是4.指数为mask.

我可以像这样得到它，但如果没有满足条件的值，它会给我一个错误.如果什么都没找到，我就拿None块.

x = df.loc[mask.cumsum().eq(1) & mask].index[0]

推荐答案

你可以切片并使用first_valid_index:

out = df[mask].first_valid_index()

# or
# out = df.loc[mask, 'a'].first_valid_index()

输出:4

输出mask = df['a'].gt(10):None

idxmax和any:

out = mask.idxmax() if mask.any() else None

或者使用python方法:

out = next((i for i,m in zip(df.index, mask) if m), None)

# alternatively (thanks @ChrisFu)
# out = next((i for i,m in mask.items() if m), None)

如果DataFrame很大，并且第一个True个索引出现得很早，这实际上比矢量化要快.

Timings

如果第一个True个提前出现(df = pd.DataFrame({'a': range(1_000_000)}) ; mask = df['a'].gt(2)):

# df.loc[mask, 'a'].first_valid_index()
5.26 ms ± 140 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# mask.idxmax() if mask.any() else None
17 µs ± 261 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

# next((i for i,m in zip(df.index, mask) if m), None)
2.69 µs ± 11.4 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

如果第一个True个出现晚(mask = df['a'].ge(999_999)):

# df.loc[mask, 'a'].first_valid_index()
426 µs ± 21.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

# mask.idxmax() if mask.any() else None
78.4 µs ± 142 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

# next((i for i,m in zip(df.index, mask) if m), None)
91.6 ms ± 1.14 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

对于1M行的输入，取决于第一个True个位置的时序图: