обратный отсчет в массиве NumPy

Question

Давайте предположим, что у меня есть ряд целочисленных значений, расположенных в массиве, подобном этому.

nan = np.nan
arr = np.array([3, nan, nan, nan, 5, nan, nan, nan, nan, nan])

nan значения должны быть заполнены обратным счетчиком от первого ненулевого значения до нуля.

[3, 2, 1, 0, 5, 4, 3, 2, 1, 0]

Answer 1

IMO, самый простой способ сделать это для панд - использовать groupby и cumcount с ascending=False:

s = pd.Series(np.cumsum(~np.isnan(arr)))
s.groupby(s).cumcount(ascending=False)

0    3
1    2
2    1
3    0
4    5
5    4
6    3
7    2
8    1
9    0
dtype: int64

Answer 2

Вот векторизованный с NumPy -

def backward_count(a):
    m = ~np.isnan(a)
    idx = np.flatnonzero(m)

    p = np.full(len(a), -1, dtype=a.dtype)
    p[idx[0]] = a[idx[0]]+idx[0]

    d = np.diff(idx)
    p[idx[1:]] = np.diff(a[m]) + d - 1
    out = p.cumsum()
    out[:idx[0]] = np.nan
    return out

Пробный прогон с более общим регистром -

In [238]: a
Out[238]: array([nan,  3., nan,  5., nan, 10., nan, nan,  4., nan, nan])

In [239]: backward_count(a)
Out[239]: array([nan,  3.,  2.,  5.,  4., 10.,  9.,  8.,  4.,  3.,  2.])

Бенчмаркинг

Настройка с масштабированием данного семпла10,000x -

In [240]: arr = np.array([3, nan, nan, nan, 5, nan, nan, nan, nan, nan])

In [241]: arr = np.tile(arr,10000)

# Pandas based one by @cs95
In [243]: %%timeit
     ...: s = pd.Series(np.cumsum(~np.isnan(arr)))
     ...: s.groupby(s).cumcount(ascending=False)
35.9 ms ± 258 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [245]: %timeit backward_count(arr)
3.04 ms ± 4.35 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 3

import pandas as pd
import numpy as np
import math

arr = pd.Series([3,np.nan,np.nan,np.nan,5,np.nan,np.nan,np.nan,np.nan,np.nan])

for i in range(len(arr)):
    # Check if each element is "NaN"
    if math.isnan(arr[i]):
        # If NaN then take the previous element and subtract 1
        arr[i] = arr[i-1]-1

# print the final array
print(arr)

Результат:

0    3.0
1    2.0
2    1.0
3    0.0
4    5.0
5    4.0
6    3.0
7    2.0
8    1.0
9    0.0
dtype: float64