Numpy中ndarray的处理函数（二）：compress()函数

mayaview

NDarray的抽象函数数目并不多，还有不少抽象函数是使用抽象函数实现的。这主要是因为抽象函数执行的操作都比较抽象，一般的用户很难将需要的操作抽象到所需要的级别。对于一般用户除了特别的情况，一般不推荐大家使用这类函数。如果你能掌握这类函数对于你操作NDarray无疑是一组强大的工具，但是即使你不关心这些函数而仅仅使用slice也可以写出挺高质量的代码。

今天我们来看看另一个很基本的NDarray的抽象函数compress。它的功能和布尔值的index是一样的。以下是这个函数的文档：

Definition: compress(condition, a, axis=None, out=None)
Docstring:
Return selected slices of an array along given axis.

When working along a given axis, a slice along that axis is returned in
`output` for each index where `condition` evaluates to True. When
working on a 1-D array, `compress` is equivalent to `extract`.

Parameters
----------
condition : 1-D array of bools
    Array that selects which entries to return. If len(condition)
    is less than the size of `a` along the given axis, then output is
    truncated to the length of the condition array.
a : array_like
    Array from which to extract a part.
axis : int, optional
    Axis along which to take slices. If None (default), work on the
    flattened array.

注意到文档里axis说的是along，而上一次讲的take是The axis over which to select values，但是遗憾的是这两个函数的选取方式是一样的。所以这个地方文档可能不太恰当，不过效果还是和take一样直观了。下面是一个例子：

In [68]: a=arange(6).reshape((3,2))
array([[0, 1],
       [2, 3],
       [4, 5]])

In [69]: b=a[:,0]>2

In [70]: b
Out[70]: array([False, False,  True], dtype=bool)

In [71]: a[b,:]
Out[71]: array([[4, 5]])

In [72]: compress(b,a)   #必须加axis，不然是在flat之后的数组上提取
Out[72]: array([2])

In [73]: compress(b,a,0)
Out[73]: array([[4, 5]])

效率上还是和take一样：

In [75]: a=random.random((1000,1000))

In [76]: b=a[:,0]>0.3

In [77]: %timeit a[b,:]
100 loops, best of 3: 8.19 ms per loop

In [78]: %timeit compress(b,a,0)
1000 loops, best of 3: 1.58 ms per loop

In [79]: array_equal(a[b,:],compress(b,a,0))
Out[79]: True


但是，compress对于flat的矩阵速度并不快，快的只是针对axis的情况（take也有相似的情况），所以需要根据情况选择不同的函数，当你不确定的时候使用你觉得容易理解的函数总是个不错的选择。

In [80]: b=a>0.4

In [81]: %timeit a[b]
100 loops, best of 3: 6.82 ms per loop

In [82]: %timeit compress(b.ravel(),a)
100 loops, best of 3: 10.5 ms per loop

In [83]: array_equal(compress(b.ravel(),a),a[b])
Out[83]: True