Python подсчитать расстояние до числа в столбце

Question

У меня есть DataFrame df, где есть целочисленный столбец 'X'

df = pd.DataFrame({'X': [7, 2, 0, 3, 4, 2, 5, 0, 3, 4]})

Мне нужно посчитать расстояние до предыдущего нуля для каждого значения в df.

, поэтомуожидаемый результат будет

op = pd.DataFrame({'dist': [0 ,0, 0, 1, 2, 3, 4, 0, 1, 2]})

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

Answer 1

Использование:

g = df['X'].eq(0).cumsum()
df['dist'] = np.where(g.eq(0), 0, df.groupby(g).cumcount())
print (df)
   X  dist
0  7     0
1  2     0
2  0     0
3  3     1
4  4     2
5  2     3
6  5     4
7  0     0
8  3     1
9  4     2

Объяснение :

1. Первое сравнение 0 с Series.eq (==)
2. Добавление кумулятивной суммы по Series.cumsum
3. Создание счетчика для групп по GroupBy.cumcount
4. Проверка первой группыg до первого 0
5. Создать новый столбец по numpy.where

print (df.assign(mask = df['X'].eq(0),
                 g = df['X'].eq(0).cumsum(),
                 counter = df.groupby(g).cumcount(),
                 mask1 = g.eq(0),
                 dist = np.where(g.eq(0), 0, df.groupby(g).cumcount())))

   X   mask  g  counter  mask1  dist
0  7  False  0        0   True     0
1  2  False  0        1   True     0
2  0   True  1        0  False     0
3  3  False  1        1  False     1
4  4  False  1        2  False     2
5  2  False  1        3  False     3
6  5  False  1        4  False     4
7  0   True  2        0  False     0
8  3  False  2        1  False     1
9  4  False  2        2  False     2