如何在引用之前的行時遍歷pandas數據框？

Question

我正在遍歷一個Python數據框，發現它非常慢。 我了解在Pandas中您嘗試對所有內容進行矢量化，但是在這種情況下，我特別需要進行迭代（或者如果可以進行矢量化，則不清楚如何執行此操作）。

邏輯很簡單：您有兩個列“ A”和“ B”，以及一個結果列“ signal”。 如果A等於1，則將signal設置為1。如果B等於1，則將signal設置為0。否則，signals與以前相同。 換句話說，列A是“開”信號，列B是“關”信號，並且“信號”表示狀態。

這是我的代碼：

def signals(indata):
    numrows = len(indata)
    data = pd.DataFrame(index= range(0,numrows))
    data['A'] = indata['A']
    data['B'] = indata['B']
    data['signal'] = 0


    for i in range(1,numrows):
        if data['A'].iloc[i] == 1:
            data['signal'].iloc[i] = 1
        elif data['B'].iloc[i] == 1:
            data['signal'].iloc[i] = 0
        else:
            data['signal'].iloc[i] = data['signal'].iloc[i-1]
    return data

輸入/輸出示例：

indata = pd.DataFrame(index = range(0,10))
indata['A'] = [0, 1, 0, 0, 0, 0, 1, 0, 0, 0]
indata['B'] = [1, 0, 0, 0, 1, 0, 0, 0, 1, 1]

signals(indata)

輸出：

    A   B   signal
0   0   1   0
1   1   0   1
2   0   0   1
3   0   0   1
4   0   1   0
5   0   0   0
6   1   0   1
7   0   0   1
8   0   1   0
9   0   1   0

這種簡單的邏輯使我的計算機需要46秒才能在具有隨機生成的數據的2000行數據幀上運行。

Answer 1

df['signal'] = df.A.groupby((df.A != df.B).cumsum()).transform('head', 1)

df
   A  B  signal
0  0  1       0
1  1  0       1
2  0  0       1
3  0  0       1
4  0  1       0
5  0  0       0
6  1  0       1
7  0  0       1
8  0  1       0
9  0  1       0

這里的邏輯涉及根據A和B之間的不等式將序列划分為組，每個組的值由A決定。

Answer 2

您完全不需要迭代就可以執行一些布爾索引

#set condition for A
indata.loc[indata.A == 1,'signal'] = 1
#set condition for B
indata.loc[indata.B == 1,'signal'] = 0
#forward fill NaN values
indata.signal.fillna(method='ffill',inplace=True)

Answer 3

解決我的問題的最簡單答案是在遍歷數據幀時不寫入數據幀。 我在numpy中創建了一個零數組，然后在該數組中執行了迭代邏輯。 然后，將數組寫入數據框的列中。

def signals3(indata):
    numrows = len(indata)
    data = pd.DataFrame(index= range(0,numrows))

    data['A'] = indata['A'] 
    data['B'] = indata['B']
    out_signal = np.zeros(numrows)

    for i in range(1,numrows):
        if data['A'].iloc[i] == 1:
            out_signal[i] = 1
        elif data['B'].iloc[i] == 1:
            out_signal[i] = 0
        else:
            out_signal[i] = out_signal[i-1]


    data['signal'] = out_signal

    return data

在2000行隨機數據的數據幀上，這只需要43毫秒，而不是46秒（快1000倍）。

我還嘗試了一個變體，在該變體中，我將數據幀的A列和B列分配給了系列，然后遍歷該系列。 這有點快（27毫秒）。 但是看來，大多數的慢是寫數據幀。

Coldspeed和djk的答案都比我的解決方案快（約4.5毫秒），但實際上，即使不是最佳選擇，我也可能會迭代整個序列。