How to find indices for sequential NaNs at the start and end of a1D NumPy array?

user_7277080 0
Viewed 163

I have a 1D NumPy array that contains both floating point numbers and NaNs. There will almost always be multiple sequential NaNs at the start and end of the array with some NaNs throughout the middle of the array. There is no way to predict in advance the location or number of NaNs present in the array.

What I am trying to do:

  • Retrieve indices of sequential NaNs occurring at the Start of the array and the End of the array only
  • Not retrieve indices for NaNs in the middle of the array
  • Avoid the use of loops/list comprehensions (if possible)

Below is a short example illustrating the type of array I am working with (actual arrays will have 1e3 to 1e9 elements).

data = np.array([NaN, NaN, NaN, 1.2, 3.4, 5.0, NaN, 3.3, 3.2, NaN, NaN, 4.0, NaN, NaN, NaN, NaN])

For this array I would like to retrieve the following indices:

idx = [0, 1, 2, 12, 13, 14, 15]
numpy python
1 Answers

Start with a mask of all the NaNs, using np.isnan:

mask = np.isnan(data)

Now notice that np.argmin always return the first instance of the minimum value:

leading_nans = mask.argmin()

You can reverse an array without copying any data using a simple slice:

trailing_nans = mask[::-1].argmin()

Either of these can be zero. Either way, you can now concatenate two ranges to get the actual indices:

nan_indices = np.concatenate((np.arange(leading_nans), np.arange(data.size - trailing_nans, data.size))

If you want a mask instead:

nan_mask = np.zeros(data.size, dtype=bool)
nan_mask[:leading_nans] = True
nan_mask[-trailing_nans:] = True
user_29887300
Related