Most efficient way to forward-fill NaN values in numpy array

后端 未结 5 1027
南笙
南笙 2020-11-28 04:25

Example Problem

As a simple example, consider the numpy array arr as defined below:

import numpy as np
arr = np.array([[5, np.nan, np.         


        
5条回答
  •  鱼传尺愫
    2020-11-28 05:18

    For those that came here looking for the backward-fill of NaN values, I modified the solution provided by Divakar above to do exactly that. The trick is that you have to do the accumulation on the reversed array using the minimum except for the maximum.

    Here is the code:

    
    
    # As provided in the answer by Divakar
    def ffill(arr):
        mask = np.isnan(arr)
        idx = np.where(~mask, np.arange(mask.shape[1]), 0)
        np.maximum.accumulate(idx, axis=1, out=idx)
        out = arr[np.arange(idx.shape[0])[:,None], idx]
        return out
    
    # My modification to do a backward-fill
    def bfill(arr):
        mask = np.isnan(arr)
        idx = np.where(~mask, np.arange(mask.shape[1]), mask.shape[1] - 1)
        idx = np.minimum.accumulate(idx[:, ::-1], axis=1)[:, ::-1]
        out = arr[np.arange(idx.shape[0])[:,None], idx]
        return out
    
    
    # Test both functions
    arr = np.array([[5, np.nan, np.nan, 7, 2],
                    [3, np.nan, 1, 8, np.nan],
                    [4, 9, 6, np.nan, np.nan]])
    print('Array:')
    print(arr)
    
    print('\nffill')
    print(ffill(arr))
    
    print('\nbfill')
    print(bfill(arr))
    
    

    Output:

    Array:
    [[ 5. nan nan  7.  2.]
     [ 3. nan  1.  8. nan]
     [ 4.  9.  6. nan nan]]
    
    ffill
    [[5. 5. 5. 7. 2.]
     [3. 3. 1. 8. 8.]
     [4. 9. 6. 6. 6.]]
    
    bfill
    [[ 5.  7.  7.  7.  2.]
     [ 3.  1.  1.  8. nan]
     [ 4.  9.  6. nan nan]]
    

    Edit: Update according to comment of MS_

提交回复
热议问题