python - Numpy: "Adding"索引和冒号

Question

我有

>>> foo = np.zeros((3,3,3))
>>> foo[1,2,1] = 1
>>> idx = 1,2

我想得到相当于

>>> foo[1,2,:]
array([ 0.,  1.,  0.])

使用idx(遍历idx)。我尝试的两种方法都没有奏效:

>>> foo[idx, :]
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  1.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])
>>> foo[((idx,)+(slice(None),))]
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  1.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])

Answer 1

使用foo[idx]。阅读以下内容，了解我是如何得出这一点的。

---

foo[1,2,:] 实际上是:

In [379]: foo[(1,2,slice(None))]
Out[379]: array([ 0.,  1.,  0.])

Python 解释器将 1,2,: 转换为这个元组，并将其传递给 foo.__getitem__ 方法。

所以我们只需要找到构造元组的正确方法。一是:

In [380]: tuple(idx)+(slice(None),)
Out[380]: (1, 2, slice(None, None, None))

导致:

In [381]: foo[tuple(idx)+(slice(None),)]
Out[381]: array([ 0.,  1.,  0.])

其实我不需要调用tuple，idx已经是一个tuple

In [386]: idx
Out[386]: (1, 2)
In [387]: idx+(slice(None),)
Out[387]: (1, 2, slice(None, None, None))

如果 idx 被初始化为列表，我将需要 tuple([1,2])。

实际上加倍，解决方案更简单。因为我们对最后一个维度使用 :，所以我们可以省略它。由于 idx 是元组，它已经索引了前 2 个维度。

In [394]: foo[idx]
Out[394]: array([ 0.,  1.,  0.])

这对调用可能会带来一些清晰度:

In [396]: foo[(1,2)]
Out[396]: array([ 0.,  1.,  0.])
In [397]: foo[[1,2]]
Out[397]: 
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  1.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])