沿着最后维度的索引numpy nd数组

有没有简单的方法来使用索引数组索引沿最后维度的numpy多维数组？例如，形状为(10, 10, 20)的数组a。假设我有一个形状为(10, 10)的索引b的数组，因此结果将为c[i, j] = a[i, j, b[i, j]]。沿着最后维度的索引numpy nd数组

我试过下面的例子：

a = np.ones((10, 10, 20)) 
b = np.tile(np.arange(10) + 10, (10, 1)) 
c = a[b]

但是，这并不工作，因为它然后尝试指数像a[b[i, j], b[i, j]]，这是不一样的a[i, j, b[i, j]]。等等。有没有一个简单的方法来做到这一点，而不诉诸于循环？

来源

2014-12-03 tiago

只是为了确保我的理解正确，你想'C [I，J] = A [1， j，b [i，j]]' – mgilson 2014-12-03 17:00:18

@mgilson是的。应该做得更清楚。 – tiago 2014-12-03 18:19:41

这通常是'c = a [np.arange（b.shape [0]），np.arange（b.shape [1]），b]'完成的，但我希望有更好的方法。 – 2014-12-03 20:31:48

有几种方法可以做到这一点。让我们先来生成一些测试数据：要解决的问题是要建立两个指标向量，其中一个是行向量和0..9使用meshgrid其他的列向量

In [1]: a = np.random.rand(10, 10, 20) 

In [2]: b = np.random.randint(20, size=(10,10)) # random integers in range 0..19

方式一：

In [3]: i1, i0 = np.meshgrid(range(10), range(10), sparse=True) 

In [4]: c = a[i0, i1, b]

这是可行的，因为i0,i1和b都将被广播到10x10矩阵。正确性快速测试：

In [5]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10)) 
Out[5]: True

另一种方法是使用choose和rollaxis：

# choose needs a sequence of length 20, so move last axis to front 
In [22]: aa = np.rollaxis(a, -1) 

In [23]: c = np.choose(b, aa) 

In [24]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10)) 
Out[24]: True

来源

2014-12-03 21:01:14

很好地完成。我喜欢。 – mgilson 2014-12-03 22:34:10

谢谢。我更喜欢你的第一个选择。 'choose'选项看起来不错，但它不够普遍，因为'choose'不适用于大于32的索引（当前的numpy错误，请参阅[关于增加'NPY_MAXARGS'的讨论]（https://github.com/numpy的/ numpy的/问题/ 4398））。根据我的测试，它也是'meshgrid'选项的两倍。 – tiago 2014-12-04 09:45:22

在这种情况下，您应该查看'mgrid'和'ogrid'等函数，这些函数是从'meshgrid'文档中引用的。但是你使用32维以上的矩阵吗？我为你感到难过，当试图理解超过2维的阵列时，我的头已经爆炸了:)。 – 2014-12-04 09:53:58

沿着最后维度的索引numpy nd数组

回答

相关问题