Python 如何将一个'numpy.ndarray'子集，其中另一个沿某个轴为max？

在

python/numpy

中，如果多维数组中的另一个相同形状的数组沿某个轴（例如第一个轴）最大，我如何将其子集

假设我有两个3*2*4数组，

a

和

b

。我想在

a

沿第一个轴具有最大值的位置获得一个包含

b

值的2*4数组

import numpy as np

np.random.seed(7)
a = np.random.rand(3*2*4).reshape((3,2,4))
b = np.random.rand(3*2*4).reshape((3,2,4))

print a
#[[[ 0.07630829  0.77991879  0.43840923  0.72346518]
#  [ 0.97798951  0.53849587  0.50112046  0.07205113]]
#
# [[ 0.26843898  0.4998825   0.67923     0.80373904]
#  [ 0.38094113  0.06593635  0.2881456   0.90959353]]
#
# [[ 0.21338535  0.45212396  0.93120602  0.02489923]
#  [ 0.60054892  0.9501295   0.23030288  0.54848992]]]

print a.argmax(axis=0) #(I would like b at these locations along axis0)
#[[1 0 2 1]
# [0 2 0 1]]

我可以做这个非常难看的手动子集：

index = zip(a.argmax(axis=0).flatten(),
            [0]*a.shape[2]+[1]*a.shape[2], # a.shape[2] = 4 here
            range(a.shape[2])+range(a.shape[2]))
# [(1, 0, 0), (0, 0, 1), (2, 0, 2), (1, 0, 3), 
#  (0, 1, 0), (2, 1, 1), (0, 1, 2), (1, 1, 3)]

这将使我能够获得我想要的结果：

b_where_a_is_max_along0 = np.array([b[i] for i in index]).reshape(2,4)

# For verification:
print a.max(axis=0) == np.array([a[i] for i in index]).reshape(2,4)
#[[ True  True  True  True]
# [ True  True  True  True]]

实现这一点的智能方法是什么？谢谢：）

使用-

样本运行-

设置输入数组

a

并沿第一个轴获取其

argmax

-

In [185]: a = np.random.randint(11,99,(3,2,4))

In [186]: idx = a.argmax(0)

In [187]: idx
Out[187]: 
array([[0, 2, 1, 2],
       [0, 1, 2, 0]])

In [188]: a
Out[188]: 
array([[[49*, 58, 13, 69],   # * are the max positions
        [94*, 28, 55, 86*]],

       [[34, 17, 57*, 50],
        [48, 73*, 22, 80]],

       [[19, 89*, 42, 71*],
        [24, 12, 66*, 82]]])

使用

b

-

In [193]: b
Out[193]: 
array([[[18*, 72, 35, 51],   # Mark * at the same positions in b
        [74*, 57, 50, 84*]], # and verify

       [[58, 92, 53*, 65],
        [51, 95*, 43, 94]],

       [[85, 23*, 13, 17*],
        [17, 64, 35*, 91]]])

In [194]: b[a.argmax(0),np.arange(2)[:,None],np.arange(4)]
Out[194]: 
array([[18, 23, 53, 17],
       [74, 95, 35, 84]])

---

您可以使用

ogrid

>>> x = np.random.random((2,3,4))
>>> x
array([[[ 0.87412737,  0.11069105,  0.86951092,  0.74895912],
        [ 0.48237622,  0.67502597,  0.11935148,  0.44133397],
        [ 0.65169681,  0.21843482,  0.52877862,  0.72662927]],

       [[ 0.48979028,  0.97103611,  0.36459645,  0.80723839],
        [ 0.90467511,  0.79118429,  0.31371856,  0.99443492],
        [ 0.96329039,  0.59534491,  0.15071331,  0.52409446]]])
>>> y = np.argmax(x, axis=1)
>>> y
array([[0, 1, 0, 0],
       [2, 0, 0, 1]])
>>> i, j = np.ogrid[:2,:4]
>>> x[i ,y, j]
array([[ 0.87412737,  0.67502597,  0.86951092,  0.74895912],
       [ 0.96329039,  0.97103611,  0.36459645,  0.99443492]])

>>> x = np.random.random((2,3,4))
>>> x
array([[[ 0.87412737,  0.11069105,  0.86951092,  0.74895912],
        [ 0.48237622,  0.67502597,  0.11935148,  0.44133397],
        [ 0.65169681,  0.21843482,  0.52877862,  0.72662927]],

       [[ 0.48979028,  0.97103611,  0.36459645,  0.80723839],
        [ 0.90467511,  0.79118429,  0.31371856,  0.99443492],
        [ 0.96329039,  0.59534491,  0.15071331,  0.52409446]]])
>>> y = np.argmax(x, axis=1)
>>> y
array([[0, 1, 0, 0],
       [2, 0, 0, 1]])
>>> i, j = np.ogrid[:2,:4]
>>> x[i ,y, j]
array([[ 0.87412737,  0.67502597,  0.86951092,  0.74895912],
       [ 0.96329039,  0.97103611,  0.36459645,  0.99443492]])