Кодирование столбца в Pandas на основе вхождения значения 0

Question

У меня есть кадр данных Pandas с таким столбцом, как

df = pd.DataFrame()
df['A'] = [1, 1, 0, 1, 1, 0]

Я хочу создать еще один столбец с такими значениями, как

[1, 1, 1, 2, 2, 2]

Идея состоит в том, чтобы начать со значения 1 и увеличивать значение, когда я получаю 1, и только если последнее значение было 0. Другими словами, если у меня есть 0, то увеличьте значение на следующем шаге.

Я использовал заявку, чтобы сделать это, как показано ниже,

k = 1

def fn(row):

    global k
    a, b = row['A'], row['x']

    if a == 1 and b == 1:
        pass
    elif a == 1 and b == 0:
        pass
    elif a == 0 and b == 1:
        k += 1
        return (k - 1)
    else:
        k += 1
        return (k - 1)

    return k

df['x'] = df['A'].shift(-1)
df['k'] = df.apply(lambda row : fn(row), axis=1)

Что действительно неэффективно. Я не могу найти более быстрый метод для этого.

Как эффективно реализовать это в Pandas .?

Answer 1

IIUC, вы хотите сосчитать вхождение 0, но смещено:

df['A'].eq(0).cumsum().shift(fill_value=0)+1

Или:

df['A'].shift().eq(0).cumsum()+1

Выход:

0    1
1    1
2    1
3    2
4    2
5    2
Name: A, dtype: int32

Answer 2

Не уверен, что скорость часть

s=df.index[df.A.eq(0)]
pd.Series(np.arange(len(s))+1,index=s).reindex(df.index,method='bfill')
Out[41]: 
0    1
1    1
2    1
3    2
4    2
5    2
dtype: int32