在Python/Numpy/Pandas中找到连续数值块的起止点

Question

我想找到一个numpy数组中相同值块的开始和结束索引，或者更好的是在pandas DataFrame中（对于二维数组，块沿着列方向，对于n维数组，块沿着变化最快的索引方向）。我只关注单一维度的块，不想把不同行中的缺失值（nan）合并在一起。

基于这个问题（在numpy数组中查找满足条件的大量连续值），我写了以下解决方案，用于查找二维数组中的np.nan：

import numpy as np
a = np.array([
        [1, np.nan, np.nan, 2],
        [np.nan, 1, np.nan, 3], 
        [np.nan, np.nan, np.nan, np.nan]
    ])

nan_mask = np.isnan(a)
start_nans_mask = np.hstack((np.resize(nan_mask[:,0],(a.shape[0],1)),
                             np.logical_and(np.logical_not(nan_mask[:,:-1]), nan_mask[:,1:])
                             ))
stop_nans_mask = np.hstack((np.logical_and(nan_mask[:,:-1], np.logical_not(nan_mask[:,1:])),
                            np.resize(nan_mask[:,-1], (a.shape[0],1))
                            ))

start_row_idx,start_col_idx = np.where(start_nans_mask)
stop_row_idx,stop_col_idx = np.where(stop_nans_mask)

这让我可以分析缺失值块的长度分布，然后再应用pd.fillna。

stop_col_idx - start_col_idx + 1
array([2, 1, 1, 4], dtype=int64)

再举一个例子和预期结果：

a = np.array([
        [1, np.nan, np.nan, 2],
        [np.nan, 1, np.nan, np.nan], 
        [np.nan, np.nan, np.nan, np.nan]
    ])

array([2, 1, 2, 4], dtype=int64)

而不是

array([2, 1, 6], dtype=int64)

我有以下几个问题：

• 有没有办法优化我的解决方案（在一次mask/where操作中找到开始和结束）？
• 在pandas中有没有更优化的解决方案？（也就是说，不只是对DataFrame的值应用mask/where）
• 当底层数组或DataFrame太大而无法放入内存时，会发生什么？

Answer 1

我把你的np.array加载到了一个数据框里：

In [26]: df
Out[26]:
    0   1   2   3
0   1 NaN NaN   2
1 NaN   1 NaN   2
2 NaN NaN NaN NaN

然后我把它转置并变成了一个序列。我觉得这和np.hstack有点像：

In [28]: s = df.T.unstack(); s
Out[28]:
0  0     1
   1   NaN
   2   NaN
   3     2
1  0   NaN
   1     1
   2   NaN
   3     2
2  0   NaN
   1   NaN
   2   NaN
   3   NaN

这个表达式创建了一个序列，里面的数字表示每个非空值对应的块，每个块的数字都增加1：

In [29]: s.notnull().astype(int).cumsum()
Out[29]:
0  0    1
   1    1
   2    1
   3    2
1  0    2
   1    3
   2    3
   3    4
2  0    4
   1    4
   2    4
   3    4

这个表达式创建了一个序列，里面每个nan（空值）都变成1，其他的都变成0：

In [31]: s.isnull().astype(int)
Out[31]:
0  0    0
   1    1
   2    1
   3    0
1  0    1
   1    0
   2    1
   3    0
2  0    1
   1    1
   2    1
   3    1

我们可以用以下方式把这两个结合起来，来得到你需要的计数：

In [32]: s.isnull().astype(int).groupby(s.notnull().astype(int).cumsum()).sum()
Out[32]:
1    2
2    1
3    1
4    4

Answer 2

下面是一个基于numpy的实现，可以处理任意维度（ndim = 2或更多）：

def get_nans_blocks_length(a):
    """
    Returns 1D length of np.nan s block in sequence depth wise (last axis).
    """
    nan_mask = np.isnan(a)
    start_nans_mask = np.concatenate((np.resize(nan_mask[...,0],a.shape[:-1]+(1,)),
                                 np.logical_and(np.logical_not(nan_mask[...,:-1]), nan_mask[...,1:])
                                 ), axis=a.ndim-1)
    stop_nans_mask = np.concatenate((np.logical_and(nan_mask[...,:-1], np.logical_not(nan_mask[...,1:])),
                                np.resize(nan_mask[...,-1], a.shape[:-1]+(1,))
                                ), axis=a.ndim-1)

    start_idxs = np.where(start_nans_mask)
    stop_idxs = np.where(stop_nans_mask)
    return stop_idxs[-1] - start_idxs[-1] + 1

这样可以：

a = np.array([
        [1, np.nan, np.nan, np.nan],
        [np.nan, 1, np.nan, 2], 
        [np.nan, np.nan, np.nan, np.nan]
    ])
get_nans_blocks_length(a)
array([3, 1, 1, 4], dtype=int64)

还有：

a = np.array([
        [[1, np.nan], [np.nan, np.nan]],
        [[np.nan, 1], [np.nan, 2]], 
        [[np.nan, np.nan], [np.nan, np.nan]]
    ])
get_nans_blocks_length(a)
array([1, 2, 1, 1, 2, 2], dtype=int64)