NumPy：将多维索引的子集传递给数组答案

【问题标题】：NumPy: pass subset of multi-dimensional indices to arrayNumPy：将多维索引的子集传递给数组
【发布时间】：2023-03-23 03:28:01
【问题描述】：

假设我有一个数组和一组（子）索引：

arr = np.array([[[0,1,2],
                 [3,4,5],
                 [6,7,8]],
                
                [[9, 10,11],
                 [12,13,14],
                 [15,16,17]]])
idx = [1,2]

并且我希望为数组的每个 dim=0 元素获得相应的idx 切片，即：

>>> arr[:, idx[0], idx[1]]
array([ 5, 14])

有没有办法在不对每个索引进行硬编码的情况下进行这种切片？比如：

ar[:, *idx]

注意，the following 是我目前的解决方法：

idx = [slice(a.shape[0]), *idx]
a[idx]

但我想知道 NumPy（或 PyTorch 或相关库）是否支持这种情况下更“优雅”/自然的多维索引语法。

【问题讨论】：

标签： python arrays numpy pytorch multi-index

【解决方案1】：

我认为您的解决方法非常接近。

你可以试试这个让它更简洁一点：

import numpy as np

arr = np.array([[[0,1,2],
                 [3,4,5],
                 [6,7,8]],

                [[9, 10,11],
                 [12,13,14],
                 [15,16,17]]])
idx = [1,2]

print(
    arr[(slice(None), *idx)]
)

这是因为: 与slice(None) 相同

【讨论】：