python - 将一个 numpy 数组取另一个数组的子集

Question

在基本层面上，我可以用另一个 numpy 数组索引一个 numpy 数组，这样我就可以返回数组的索引，这样:

a = [1,2,3,4,5,6]

b = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9]

索引 0.6 可以通过以下方式找到:

c = a[b==0.6]

但是，现在我有了 3D 数组，但我无法弄清楚我需要什么。

我有 3 个数组:

A = [[21,22,23....48,49,50]] # An index over the range 20-50 with shape (1,30)

B = [[0.1,0.6,0.5,0.4,0.8...0.7,0.2,0.4],
     ..................................
     [0.5,0.2,0.7,0.1,0.5...0.8,0.9,0.3]] # This is my data with shape (40000, 30)

C = [[0.8],........[0.9]] # Maximum values from each array in B with shape (40000,1)

我想通过索引数据 (B) 中每个数组中的最大值 (C) 来了解位置(从 A)

我已经尝试过:

D = A[B==C]

但我不断收到错误:

IndexError: index 1 is out of bounds for axis 0 with size 1

单独我可以得到:

B==C # prints as arrays of True or False

但我无法从 A 检索索引位置。

感谢任何帮助!

Answer 1

这是您正在寻找的吗？使用 argmax 函数获取每行最大值所在的索引，并使用索引获取 A 中的相应值。

In [16]: x = np.random.random((20, 30))
In [16]: max_inds = x.argmax(axis=1)

In [17]: max_inds.shape
Out[17]: (20,)

In [18]: A = np.arange(x.shape[1])

In [19]: A.shape
Out[19]: (30,)

In [20]: A[max_inds]
Out[20]: 
array([20,  5, 27, 19, 27, 21, 18, 25, 10, 24, 16, 21,  6,  7, 27, 17, 24,
        8, 27,  8])