矢量化/加速循环

Question

我有一个 pandas DataFrame "data" 并且想要执行以下操作：

import pandas as pd

data = pd.DataFrame({"col": [0,0,0,1,0,0,0,0.1,5,0,0,0],
                    "t": [0,0,0,0,0,0,0,0,0,0,0,0]})

for i in range(1,len(data)):
        if(data["col"][i]==0):
            data["t"][i] = data["t"][i-1]+1
        else:
            data["t"][i] = 0

它工作正常，但不幸的是它真的很慢。有没有可能加快这段代码的速度？例如矢量化？

预期输出

df = pd.DataFrame({"col": [0,0,0,1,0,0,0,0.1,5,0,0,0],
                     "t": [0,1,2,0,1,2,3,0,0,1,2,3]})

    col  t
0   0.0  0
1   0.0  1
2   0.0  2
3   1.0  0
4   0.0  1
5   0.0  2
6   0.0  3
7   0.1  0
8   5.0  0
9   0.0  1
10  0.0  2
11  0.0  3

Answer 1

IIUC，你需要一个简单的np.where声明

np.where(df.col.eq(0), df.t.shift() + 1, 0)

编辑：

由于缺乏示例，真的很难理解 OP 想要什么，但是 IIUC，以下应该可行：

ans2 = df.groupby((df.col.eq(0).shift() != df.col.eq(0)).cumsum()).t.cumcount() + 1

df['ans2'] = np.where(df.col.eq(0), ans2, 0)

矢量化/加速循环

Vectorization / speed up for-loop

python

vectorization

pandas