np.argmax 那个 returns 元组

Question

我有一个形状为 (n, m, s) 的矩阵 A。在第 0 个轴的每个位置，我需要对应于 (m, s) 形数组中最大值的位置。

例如：

np.random.seed(1)
A = np.random.randint(0, 10, size=[10, 3, 3])

A[0] 是：

array([[5, 8, 9],
       [5, 0, 0],
       [1, 7, 6]])

我要获取(0, 2)，即这里9的位置

我愿意做 aa = A.argmax()，这样 aa.shape = (10, 2)，并且 aa[0] = [0, 2]

我怎样才能做到这一点？

Answer 1

使用 np.unravel_index 和列表理解：

out = [np.unravel_index(np.argmax(block, axis=None), block.shape) for block in array]

其中 block 将是每个回合中的 3x3 (m x s) 形状的数组。

这给出了一个包含 10 (n) 个条目的列表：

[(0, 2), (0, 0), (0, 1), (0, 1), (2, 0), (1, 2), (0, 0), (1, 1), (1, 0), (1, 2)]

您可以将其转换为 numpy 数组（所需形状 (n, 2)）：

aa = np.array(out)

获得：

array([[0, 2],
       [0, 0],
       [0, 1],
       [0, 1],
       [2, 0],
       [1, 2],
       [0, 0],
       [1, 1],
       [1, 0],
       [1, 2]], dtype=int64)

Answer 2

其中一个 vectorized-version 解决方案由 Mad Physicist 给出推荐用于速度和简单的矢量操作，但您可以使用基本的 python 内置操作找到解决方案。

逻辑：

np.argmax return flattened 中的最大值 axis default argument 是 None 这意味着它通过将矩阵作为向量来对矩阵进行运算。
1. // integer operation 将导致 row-number 并且
2. % module operation 将导致 column-number

result = [(np.argmax(i)//3, (np.argmax(i)%3)) for i in array]

对于Python 3.8+使用:=运算符进行单次计算：

result = [(y//3, y%3) for i in array if (y := np.argmax(i))]

代码：

import numpy as np

np.random.seed(1)
array = np.random.randint(0,10,size=[10, 3, 3])

result = [(np.argmax(i)//3, (np.argmax(i)%3)) for i in array]
print(result)

输出：

[(0, 2), (0, 0), (0, 1), (0, 1), (2, 0), (1, 2), (0, 0), (1, 1), (1, 0), (1, 2)]

Answer 3

通常运行 argmax 数组的部分展平版本比使用理解要快得多：

ind = A.reshape(A.shape[0], -1).argmax(axis=-1)

现在可以申请unravel_index获取二维索引：

coord = np.unravel_index(ind, A.shape[1:])

coord是一对数组，每个轴一个。您可以使用常见的 python 成语转换为元组列表：

result = list(zip(*coord))

更好的方法是将数组堆叠在一起形成一个 Nx2 数组。对于所有意图和目的，这将表现为一系列对：

result = np.stack(coord, axis=-1)

必填一行：

result = np.stack(np.unravel_index(A.reshape(A.shape[0], -1).argmax(axis=-1), A.shape[1:]), axis=-1)

np.argmax 那个 returns 元组

np.argmax that returns tuple

python

numpy

argmax