arrays - Matlab中沿维度子集的多维数组的Argmax

Question

比如说，Y 是一个 7 维数组，我需要一种有效的方法来沿最后 3 维最大化它，这将在 GPU 上工作。因此，我需要一个最大值为 Y 的 4 维数组和三个 4 维数组，这些数组在最后三个维度中具有这些值的索引。我可以

[Y7, X7] = max(Y , [], 7);
[Y6, X6] = max(Y7, [], 6);
[Y5, X5] = max(Y6, [], 5);

然后我已经找到了沿第 5 维 (X5) 的值 (Y5) 和索引。但我仍然需要第 6 维和第 7 维的索引。

Answer 1

这是一种方法。让N表示要最大化的维度数。

1. 重塑以将最后的尺寸Y折叠成一个。N
2. 沿折叠尺寸最大化。这将 argmax 作为这些维度上的线性索引。
3. 将线性索引展开为子索引N，每个维度一个。

以下代码适用于任意数量的维度（不一定7和3您的示例一样）。为了实现这一点，它通常处理 的大小，Y并使用从元胞数组获得的逗号分隔列表N从sub2ind.

Y = rand(2,3,2,3,2,3,2); % example 7-dimensional array
N = 3; % last dimensions along which to maximize
D = ndims(Y);
sz = size(Y);
[~, ind] = max(reshape(Y, [sz(1:D-N) prod(sz(D-N+1:end))]), [], D-N+1);
sub = cell(1,N);
[sub{:}] = ind2sub(sz(D-N+1:D), ind);

作为检查，在运行上述代码后，观察示例Y(2,3,1,2,:)（为方便起见显示为行向量）：

>> reshape(Y(2,3,1,2,:), 1, [])
ans =
    0.5621    0.4352    0.3672    0.9011    0.0332    0.5044    0.3416    0.6996    0.0610    0.2638    0.5586    0.3766

最大值被视为0.9011，它出现在第4th 位置（“位置”沿N=3折叠尺寸定义）。实际上，

>> ind(2,3,1,2)
ans =
     4
>> Y(2,3,1,2,ind(2,3,1,2))
ans =
    0.9011

或者，就子指数而言N=3，

>> Y(2,3,1,2,sub{1}(2,3,1,2),sub{2}(2,3,1,2),sub{3}(2,3,1,2))
ans =
    0.9011