计算有多少连续行满足条件 pandas

Question

我有一张这样的桌子：

import pandas as pd
df = pd.DataFrame({
    "day": [1, 2, 3, 4, 5, 6],
    "tmin": [-2, -3, -1, -4, -4, -2]
})

我想创建一个这样的列：

df['days_under_0_until_now'] = [1, 2, 3, 4, 5, 6]
df['days_under_-2_until_now'] = [1, 2, 0, 1, 2, 3]
df['days_under_-3_until_now'] = [0, 1, 0, 1, 2, 0]

所以days_under_X_until_now表示到现在为止连续多少天tmin低于或等于X

我想避免使用循环来执行此操作，因为数据量很大。 有其他方法吗？

Answer 1

为了提高性能，避免使用groupby比较要列出的列的值，然后使用此解决方案计算连续的True s：

vals = [0,-2,-3]

arr = df['tmin'].to_numpy()[:, None] <= np.array(vals)[ None, :]
cols = [f'days_under_{v}_until_now' for v in vals]
df1 = pd.DataFrame(arr, columns=cols, index=df.index)

b = df1.cumsum()
df = df.join(b.sub(b.mask(df1).ffill().fillna(0)).astype(int))
print (df)
   day  tmin  days_under_0_until_now  days_under_-2_until_now  \
0    1    -2                       1                        1   
1    2    -3                       2                        2   
2    3    -1                       3                        0   
3    4    -4                       4                        1   
4    5    -4                       5                        2   
5    6    -2                       6                        3   

   days_under_-3_until_now  
0                        0  
1                        1  
2                        0  
3                        1  
4                        2  
5                        0