Numpy:沿任意维找到非零值

时间:2016-08-26 09:08:16

标签: arrays numpy matrix iteration numpy-broadcasting

我似乎无法在Numpy中解决这个问题:我有一个矩阵,具有任意数量的维度,以任意方式排序。在这个矩阵中,总有一个我感兴趣的维度(正如我所说,这个维度的位置并不总是相同的)。现在,我想在此维度中找到第一个非零值。实际上,我需要该值的索引来对值本身执行一些操作。 例如:如果我的矩阵an x m x p且我感兴趣的维度为1,我会做类似的事情:

for ii in xrange(a.shape[0]):
   for kk in xrange(a.shape[2]):
      myview = np.squeeze(a[ii, :, kk])
      firsti = np.nonzero(myview)[0][0]
      myview[firsti] = dostuff

除了性能方面的考虑,我真的不知道如何使用不同数量的维度做这个,并且让我感兴趣的维度处于任意位置。

1 个答案:

答案 0 :(得分:0)

您可以为了您的目的滥用np.argmax。在这里,您可以指定您感兴趣的axis,其中0位于列的位置,1位于行中,依此类推。您只需要一个数组,其中包含非零的所有元素的相同值。您可以通过执行a != 0来实现这一目标,因为这将包含所有零元素的False(意味着0)和所有零元素的True(意为1)非零元素。现在np.argmax(a != 0, axis=1)将为您提供1轴上的第一个非零元素。

例如:

import numpy as np

a = np.array([[0, 1, 4],[1, 0, 2],[0, 0, 1]])
# a = [[0, 1, 4],
#      [1, 0, 2],
#      [0, 0, 1]]

print(np.argmax(a!=0, axis=0))
# >>> array([1, 0, 0]) -> along columns

print(np.argmax(a!=0, axis=1))
# >>> array([1, 0, 2]) -> along rows

这也适用于更高维度,但输出不那么有启发性,因为很难想象。