numpy ndarray构建符合条件的索引数组——使用np.stack和np.where

  首先介绍一下np.stack和np.where两个方法：
  

1 np.stack()

  np.stack()函数定义为numpy.stack(arrays, axis=0)，用于将多个具有相同形状的array按某个维度进行堆叠，默认是在0维度上（可以理解为在哪个维度上就是增加哪个维度）。更为详细的介绍可参考这里，如下一个简单例子：

>>> import numpy as np
>>> x = np.arange(9).reshape([3,3])
>>> y = np.ones([3,3])
>>> z = np.stack((x,y))
>>> z
array([[[0., 1., 2.],
        [3., 4., 5.],
        [6., 7., 8.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]])
>>> z.shape
(2, 3, 3)

  

2 np.where()

  np.where()函数定义为numpy.where(condition[,x,y])，用于进行条件判断，满足条件的位置值为x，不满足的为y。不给x和y这两个参数的话，则是返回一系列坐标值的索引array。更为详细的介绍可参考这里，接着上面的结果，如下一个简单例子：

>>> np.where(z>1)
(array([0, 0, 0, 0, 0, 0, 0]), array([0, 1, 1, 1, 2, 2, 2]), array([2, 0, 1, 2, 0, 1, 2]))
>>> np.where(z>1,3,1)
array([[[1, 1, 3],
        [3, 3, 3],
        [3, 3, 3]],

       [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]]])

  

3 条件筛选

  其实根据上述两个小例子，我们已经可以实现对条件进行筛选的功能了，如下所示：

>>> np.stack(np.where(z>1))
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 2, 2, 2],
       [2, 0, 1, 2, 0, 1, 2]])

  我们先输出z的值对比看看：

>>> z
array([[[0., 1., 2.],
        [3., 4., 5.],
        [6., 7., 8.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]])

  对比上述两个结果，我们以列的方式来看stack输出的结果。共有7列，说明z中有7个数据是满足条件z>1的，而每一列则是这些满足条件的数据所在的坐标，比如位于[0,0,2]的是2，位于[0,1,0]的是3，等等。
  根据这个嵌套用法，我们就可以构建出对一个ndarray数据进行筛选，来获取其中满足我们所需条件的数据的坐标，抑或是直接使用where的xy参数来生成一个想要的新array。
  

4 直接抽取元素

  当然，如果要直接把符合条件的数据抽取出来，而不是要其具体坐标值，那就更简单了，直接上例子：

# 可以直接抽取
>>> z[z>1]
array([2., 3., 4., 5., 6., 7., 8.])
# 也可以先将符合条件的数据存成mask再抽取
>>> mask = z>1
>>> mask
array([[[False, False,  True],
        [ True,  True,  True],
        [ True,  True,  True]],

       [[False, False, False],
        [False, False, False],
        [False, False, False]]])
>>> z[mask]
array([2., 3., 4., 5., 6., 7., 8.])

  总而言之就是，感觉numpy的张量处理操作很灵活多变，如果足够熟悉的话，很多复杂的操作都可以通过简单的几行甚至一行代码搞定。就，多看多学多记录咯。