如何在大熊猫数据框中加快迭代功能？

Question

I am quite new to pandas and I have a pandas dataframe of about 500,000 rows filled with numbers. I am using python 2.x and am currently defining and calling the method shown below on it. It sets a predicted value to be equal to the corresponding value in series 'B', if two adjacent values in series 'A' are the same. However, it is running extremely slowly, about 5 rows are outputted per second and I want to find a way accomplish the same result more quickly.

def myModel(df):

    A_series = df['A']
    B_series = df['B']
    seriesLength = A_series.size

    # Make a new empty column in the dataframe to hold the predicted values
    df['predicted_series'] = np.nan

    # Make a new empty column to store whether or not
    # prediction matches predicted matches B
    df['wrong_prediction'] = np.nan
    prev_B = B_series[0]
    for x in range(1, seriesLength):

        prev_A = A_series[x-1]  
        prev_B = B_series[x-1]
        #set the predicted value to equal B if A has two equal values in a row
        if A_series[x] == prev_A:
            if df['predicted_series'][x] > 0:
                 df['predicted_series'][x] = df[predicted_series'][x-1]
            else:
                 df['predicted_series'][x] = B_series[x-1]

Is there a way to vectorize this or to just make it run faster? Under the current circumstances, it is projected to take many hours. Should it really be taking this long? It doesn't seem like 500,000 rows should be giving my program that much problem.

Answer 1

像您描述的那样，这样的事情应该起作用：

df['predicted_series'] = np.where(A_series.shift() == A_series, B_series, df['predicted_series'])

Answer 2

df.loc[df.A.diff() == 0, 'predicted_series'] = df.B  

This will get rid of the for loop and set predicted_series to the value of B when A is equal to previous A.

edit:

per your comment, change your initialization of predicted_series to be all NAN and then front fill the values:

df['predicted_series'] = np.nan
df.loc[df.A.diff() == 0, 'predicted_series'] = df.B
df.predicted_series = df.predicted_series.fillna(method='ffill')

For fastest speed modifying ayhans answer a bit will perform best:

df['predicted_series'] = np.where(df.A.shift() == df.A, df.B, df['predicted_series'].shift())

That will give you your forward filled values and run faster than my original recommendation

Answer 3

解

df.loc[df.A == df.A.shift()] = df.B.shift()