python - 矢量化/加速 for 循环

Question

我有一个 pandas DataFrame“数据”并且想要执行以下操作:

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)