使用 2D 矩阵作为 numpy 中 3D 矩阵的索引？答案

【问题标题】：Use 2D matrix as indexes for a 3D matrix in numpy?使用 2D 矩阵作为 numpy 中 3D 矩阵的索引？
【发布时间】：2018-01-02 06:47:16
【问题描述】：

假设我有一个形状为 2x3x3 的数组，它是一个 3D 矩阵。我还有一个形状为 3x3 的 2D 矩阵，我想将其用作沿第一个轴的 3D 矩阵的索引。示例如下。

示例运行：

>>> np.random.randint(0,2,(3,3)) # index
array([[0, 1, 0],
       [1, 0, 1],
       [1, 0, 0]])

>> np.random.randint(0,9,(2,3,3)) # 3D matrix
array([[[4, 4, 5],
        [2, 6, 7],
        [2, 6, 2]],

       [[4, 0, 0],
        [2, 7, 4],
        [4, 4, 0]]])
>>> np.array([[4,0,5],[2,6,4],[4,6,2]]) # result
array([[4, 0, 5],
       [2, 6, 4],
       [4, 6, 2]])

【问题讨论】：

2D 阵列的形状与 3D 阵列的形状有何关系？
它具有相同的 2D 形状，但只是 2D。所以 3D 数组和 2D 数组的形状 NxN 相同，但是有更多的“层”。
您的 2 个不同的矩阵是否像示例中一样重复数字？否则，您将需要提供有关您希望输出外观的更多信息。例如，0 和 1 对您意味着什么？
目前范围太广。请进一步说明输出在做什么，我可以想到至少三种不同的方法来使用这三个二维数组来获得结果
@MatthewMage 使用随机数而不是那些重复的数字。这势必会让人感到困惑。

标签： python numpy matrix

【解决方案1】：

您似乎使用2D 数组作为索引数组，使用3D 数组来选择值。因此，您可以使用 NumPy 的 advanced-indexing -

# a : 2D array of indices, b : 3D array from where values are to be picked up
m,n = a.shape
I,J = np.ogrid[:m,:n]
out = b[a, I, J] # or b[a, np.arange(m)[:,None],np.arange(n)]

如果您打算使用a 来索引最后一个轴，只需将a 移动到那里：b[I, J, a]。

示例运行 -

>>> np.random.seed(1234)
>>> a = np.random.randint(0,2,(3,3))
>>> b = np.random.randint(11,99,(2,3,3))
>>> a  # Index array
array([[1, 1, 0],
       [1, 0, 0],
       [0, 1, 1]])
>>> b  # values array
array([[[60, 34, 37],
        [41, 54, 41],
        [37, 69, 80]],

       [[91, 84, 58],
        [61, 87, 48],
        [45, 49, 78]]])
>>> m,n = a.shape
>>> I,J = np.ogrid[:m,:n]
>>> out = b[a, I, J]
>>> out
array([[91, 84, 37],
       [61, 54, 41],
       [37, 49, 78]])

【讨论】：

【解决方案2】：

如果您的矩阵变得比 3x3 大得多，那么np.ogrid 所涉及的内存是一个问题，并且如果您的索引仍然是二进制，您也可以这样做：

np.where(a, b[1], b[0])

但除了那个角落案例（或者如果你喜欢打代码打高尔夫球的单线），另一个答案可能更好。

【讨论】：

【解决方案3】：

有一个现成的 numpy 函数：np.choose。它还带有一些方便的广播选项。

import numpy as np    
cube = np.arange(18).reshape((2,3,3))
sel = np.array([[1, 0, 1], [0, 1, 1], [0,1,0]])

the_selection = np.choose(sel, cube)


>>>the_selection
array([[ 9,  1, 11],
       [ 3, 13, 14],
       [ 6, 16,  8]])

此方法适用于任何 3D 数组。

【讨论】：

你能编辑你的代码来处理形状不匹配的问题吗？此方法不适用于 OP 的问题。
如果其他人发现这篇文章，np.choose 可以很好地解决 OP 的问题，如此处所示，为了清楚起见，使用范围而不是随机 3D 数组。但是，在二进制 2D 数组和两层 3D 数组的情况下，Daniel F 使用 np.where 的智能答案更快。
不，它不起作用，你在说什么范围？ OP 有一个形状为 (2,3,3) 的变量。如果您尝试使用非多维数据集运行代码，它会中断。
我已经更新了答案以生成具有与 OP 一样的精确尺寸的 3D 数组，这也可以。 np.arange 生成一系列值，以便您自己查看从哪个层选择了值。 OP 的标题表明他对一般情况感兴趣，但使用了第一个轴 = 2 的 3D 数组的示例（“假设我有一个形状为 2x3x3 的数组”）。
这是最好的答案。我不确定choose与ogrid的性能，但清洁度是无与伦比的