python - 在 numpy 中 take 的多维等价物是什么

Question

我有这段代码

def build_tree_base(blocks, x, y, z):
   indicies = [
        (x  ,z  ,y  ),
        (x  ,z+1,y  ),
        (x  ,z  ,y+1),
        (x  ,z+1,y+1),
        (x+1,z  ,y  ),
        (x+1,z+1,y  ),
        (x+1,z  ,y+1),
        (x+1,z+1,y+1),
    ]
    children = [blocks[i] for i in indicies]
    return Node(children=children)

其中 blocks 是一个 3 维 numpy 数组。

我想做的是用 numpy.take 之类的东西替换列表理解，但是 take 似乎只处理单维索引。是否有类似 take 的东西可以用于多维索引？

此外，我知道您可以通过转置、切片然后 reshape 来完成此操作，但这很慢，所以我正在寻找更好的选择。

Answer 1

Numpy 索引使这很容易......你应该能够做到这样的事情:

def build_tree_base(blocks, x, y, z):
    idx = [x, x, x, x, x+1, x+1, x+1, x+1]
    idz = [z, z+1, z, z+1, z, z+1, z, z+1]
    idy = [y, y, y+1, y+1, y, y, y+1, y+1]
    children = blocks[idx, idz, idy]
    return Node(children=children)

编辑:我应该指出这个(或任何其他 "fancy" indexing )将返回一个副本，而不是原始数组的 View ...