r - Replace NA with previous value with limit

Question

I would like to replace up to n NA values in vector with latest non-NA value.

For example, if:

a <- c(1,NA,NA,NA,NA,NA,2,NA,1,NA,NA,NA)
n <- 2

I would like to obtain:

c(1,1,1,NA,NA,NA,2,2,1,1,1,NA)

n is maximum number of NA values that can be replaced by given element).

I know na.locf() function, but I don't know how to set the limit n. Is it possible to do it?

Answer 1

Here's an option using na.locf and rle

library(zoo)
r <- rle(is.na(a))
a <- na.locf(a)
is.na(a) <- sequence(r$lengths) > n & rep(r$values, r$lengths)
a
# [1]  1  1  1 NA NA NA  2  2  1  1  1 NA

So here I first computed the run lengths of elements in a (including the NA entries), then replaced all NA's using na.locf and finally turned those elements back to NA's where the run lengths were greater than n and the elements were NA.