numpy：将由NaN分隔的1D数组拆分为多个块的列表

Question

我有一个numpy数组，里面只有一些值是有效的，其余的都是nan（表示没有值）。举个例子：

[nan,nan, 1 , 2 , 3 , nan, nan, 10, 11 , nan, nan, nan, 23, 1, nan, 7, 8]

我想把这个数组分成几个小块，每个小块里都包含有效的数据。最终的结果应该是这样的：

[[1,2,3], [10,11], [23,1], [7,8]]

我通过遍历这个数组，检查哪些值是有效的，然后生成（开始，结束）索引来实现这个功能。

不过……这个方法实在是太慢了……

你有没有更好的办法呢？

Answer 1

我会使用 itertools.groupby 这个工具——它可能会稍微快一点：

from numpy import NaN as nan
import numpy as np
a = np.array([nan,nan, 1 , 2 , 3 , nan, nan, 10, 11 , nan, nan, nan, 23, 1, nan, 7, 8])
from itertools import groupby
result = [list(v) for k,v in groupby(a,np.isfinite) if k]
print result #[[1.0, 2.0, 3.0], [10.0, 11.0], [23.0, 1.0], [7.0, 8.0]]

Answer 2

这里还有一种可能性：

import numpy as np
nan = np.nan

def using_clump(a):
    return [a[s] for s in np.ma.clump_unmasked(np.ma.masked_invalid(a))]

x = [nan,nan, 1 , 2 , 3 , nan, nan, 10, 11 , nan, nan, nan, 23, 1, nan, 7, 8]

In [56]: using_clump(x)
Out[56]: 
[array([ 1.,  2.,  3.]),
 array([ 10.,  11.]),
 array([ 23.,   1.]),
 array([ 7.,  8.])]

---

下面是一些关于使用 using_clump 和 using_groupby 的性能对比：

import itertools as IT
groupby = IT.groupby
def using_groupby(a):
    return [list(v) for k,v in groupby(a,np.isfinite) if k]

---

In [58]: %timeit using_clump(x)
10000 loops, best of 3: 37.3 us per loop

In [59]: %timeit using_groupby(x)
10000 loops, best of 3: 53.1 us per loop

对于更大的数组，性能表现甚至更好：

In [9]: x = x*1000
In [12]: %timeit using_clump(x)
100 loops, best of 3: 5.69 ms per loop

In [13]: %timeit using_groupby(x)
10 loops, best of 3: 60 ms per loop