Filling missing values with lagged observations

lagged.py

import numpy as np


def lag(x, n=0):
    observed = np.isnan(x) == False
    observed[0] = True
    offsets = np.flatnonzero(observed)
    return offsets[np.maximum(np.cumsum(observed) - 1 - n, 0)]


if __name__ == '__main__':
    x = np.nan * np.ones(10)
    x[2] = 3
    x[3] = 5
    x[5] = 7

    print(x)
    for l in range(3):
        ii = lag(x, l)
        print(ii, x[ii])

% python3 -i lag.py
[nan nan  3.  5. nan  7. nan nan nan nan]
[0 0 2 3 3 5 5 5 5 5] [nan nan  3.  5.  5.  7.  7.  7.  7.  7.]
[0 0 0 2 2 3 3 3 3 3] [nan nan nan  3.  3.  5.  5.  5.  5.  5.]
[0 0 0 0 0 2 2 2 2 2] [nan nan nan nan nan  3.  3.  3.  3.  3.]