Cumsum reset при NaN

Если у меня есть pandas.core.series.Series с именем ts из 1 или NaN, как это:

3382   NaN
3381   NaN
...
3369   NaN
3368   NaN
...
15     1
10   NaN
11     1
12     1
13     1
9    NaN
8    NaN
7    NaN
6    NaN
3    NaN
4      1
5      1
2    NaN
1    NaN
0    NaN

Я хотел бы рассчитать cumsum этой серии, но он должен быть reset (установлен на ноль) в местоположении NaN, как показано ниже:

3382   0
3381   0
...
3369   0
3368   0
...
15     1
10     0
11     1
12     2
13     3
9      0
8      0
7      0
6      0
3      0
4      1
5      2
2      0
1      0
0      0

В идеале я хотел бы иметь векторизованное решение!

Я когда-либо видел аналогичный вопрос с Matlab: Matlab cumsum reset в NaN?

но я не знаю, как перевести эту строку d = diff([0 c(n)]);

ОТВЕТЫ

Ответ 1

Простой перевод Numy вашего кода Matlab выглядит следующим образом:

import numpy as np

v = np.array([1., 1., 1., np.nan, 1., 1., 1., 1., np.nan, 1.])
n = np.isnan(v)
a = ~n
c = np.cumsum(a)
d = np.diff(np.concatenate(([0.], c[n])))
v[n] = -d
np.cumsum(v)

Выполнение этого кода возвращает результат array([ 1., 2., 3., 0., 1., 2., 3., 4., 0., 1.]). Это решение будет только таким же правильным, как и оригинальное, но, возможно, оно поможет вам придумать что-то лучшее, если этого недостаточно для ваших целей.

Ответ 2

Здесь немного более pandas -нихий способ сделать это:

v = Series([1, 1, 1, nan, 1, 1, 1, 1, nan, 1], dtype=float)
n = v.isnull()
a = ~n
c = a.cumsum()
index = c[n].index  # need the index for reconstruction after the np.diff
d = Series(np.diff(np.hstack(([0.], c[n]))), index=index)
v[n] = -d
result = v.cumsum()

Обратите внимание, что для любого из них требуется, чтобы вы использовали pandas по крайней мере на 9da899b или новее. Если вы этого не сделаете, вы можете направить bool dtype на int64 или float64 dtype:

v = Series([1, 1, 1, nan, 1, 1, 1, 1, nan, 1], dtype=float)
n = v.isnull()
a = ~n
c = a.astype(float).cumsum()
index = c[n].index  # need the index for reconstruction after the np.diff
d = Series(np.diff(np.hstack(([0.], c[n]))), index=index)
v[n] = -d
result = v.cumsum()

Ответ 3

Еще более pandas -нический способ сделать это:

v = pd.Series([1., 3., 1., np.nan, 1., 1., 1., 1., np.nan, 1.])
cumsum = v.cumsum().fillna(method='pad')
reset = -cumsum[v.isnull()].diff().fillna(cumsum)
result = v.where(v.notnull(), reset).cumsum()

В отличие от кода matlab, это также работает для значений, отличных от 1.

Ответ 4

Если вы можете принять аналогичную логическую серию b, попробуйте

(b.cumsum() - b.cumsum().where(~b).fillna(method='pad').fillna(0)).astype(int)

Начиная с вашей серии ts, либо b = (ts == 1), либо b = ~ts.isnull().