Как я могу векторизовать мою функцию, чтобы ускорить работу на моем фрейме данных?

Question

У меня есть фрейм данных (ниже, т. Е. Членство), в одном поле (A) есть несколько строк со значением в отсортированном виде. Существует также новое поле (новое), которое в начале процесса является копией поля C. Я хотел бы сделать следующее: если предыдущая строка в A совпадает с текущей строкой в ​​A, и если либо текущая строка new, либо предыдущая строка new равна 1 присвоить 1 текущему new. В конце, в последнем случае повторных значений A, new будет 1 или 0 в зависимости от условий в функции, а предыдущие значения, где повторяется A, будут иметь new быть 0. Я могу сделать это с помощью функции ниже.

membership = pd.DataFrame.from_dict(dict([('A', ['20000000460', '20000000460', '20000000460','20000000460','20000000459','20000000461','20000000461','20000000462','20000000464','20000000464','20000000464','20000000464','20000000465','20000000465','20000000466']), ('B', [4,0, 5,0, 6,0,2,5,6,7,4,3,2,7,9]), ('C', [1,1,0,0,0,1,0,1,1,1,0,0,0,0,1])]))

def members(df, field):
df[field] = df.C
print(field)
for i in range(1, df.shape[0]):
    if (df.loc[i, 'A'] == df.loc[i-1, 'A']) and\
    (df.loc[i-1, field] == 1 or df.loc[i, field] == 1):
        df.loc[i, field] = 1
        df.loc[i-1, field] = 0

Результаты этой функции на фрейме данных в этом введите описание изображения здесь

Проблема в том, что у меня очень большой набор данных, и эта функция запускается очень медленно. Как я могу улучшить код, чтобы сделать его быстрее? Я знаю, если я смогу векторизовать эту функцию в пандах, время значительно улучшится. Как я могу векторизовать эту функцию?

Answer 1

IIUC, позвольте мне объяснить небольшую логику и посмотреть, соответствует ли это.

Если в какой-либо группе A значение C равно 1, присвойте последним записям в этой группе значение 1 для столбца «new».

membership['new'] = membership.groupby('A')['C']\
                              .transform(lambda x: np.where(x.index == x.index[-1], 
                                                            x.max(), 0))

Выход:

              A  B  C  new
0   20000000460  4  1    0
1   20000000460  0  1    0
2   20000000460  5  0    0
3   20000000460  0  0    1
4   20000000459  6  0    0
5   20000000461  0  1    0
6   20000000461  2  0    1
7   20000000462  5  1    1
8   20000000464  6  1    0
9   20000000464  7  1    0
10  20000000464  4  0    0
11  20000000464  3  0    1
12  20000000465  2  0    0
13  20000000465  7  0    0
14  20000000466  9  1    1