Numpy: 如何根据矩阵的非零索引提取元素到子矩阵中

Question

我有一些非常稀疏的矩阵，所以我想提取出一个最小的矩形区域，这个区域里包含非零值。我知道使用numpy.nonzero(a)可以得到非零元素的索引，但我该如何利用这些索引来提取出包含这些元素的子矩阵呢？

举个例子，这就是我想要的结果：

>>> test
array([[0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 1, 0],
       [0, 0, 1, 1, 0, 0]])
>>> np.nonzero(test)
(array([1, 1, 1, 1, 2, 2]), array([1, 2, 3, 4, 2, 3]))
>>> submatrix(test)
array([[1, 1, 1, 1],
       [0, 1, 1, 0]])

有没有人知道在numpy中简单的方法来做到这一点？谢谢。

Answer 1

看起来你想找到一个矩阵中包含所有非零元素的最小区域。如果是这样的话，这里有一种方法：

import numpy as np

def submatrix(arr):
    x, y = np.nonzero(arr)
    # Using the smallest and largest x and y indices of nonzero elements, 
    # we can find the desired rectangular bounds.  
    # And don't forget to add 1 to the top bound to avoid the fencepost problem.
    return arr[x.min():x.max()+1, y.min():y.max()+1]

test = np.array([[0, 0, 0, 0, 0, 0],
                 [0, 1, 1, 1, 1, 0],
                 [0, 0, 1, 1, 0, 0]])

print submatrix(test)

# Result:  
# [[1 1 1 1]
#  [0 1 1 0]]