使用np.maximum.reduceat对二维数组进行归约

Question

我有一个二维数组，这个数组里存了一对浮点数数据和它们的索引（形状是(3, k)的浮点数组）。

我想要从data[0]中找到每个索引在data[2]里的最大值，同时也想从data[1]中获取相同索引的数据。

下面是个例子：

data = np.array([[ 0.45114132,  0.31522008,  0.66176217,  0.45114132,  0.26872137],
       [11.        ,  6.        , 10.        ,  4.        ,  8.        ],
       [ 0.        ,  0.        ,  0.        ,  1.        ,  1.        ]])

# Expected Output:
array([[0.66176217, 0.45114132], [10., 4], [0, 1])

我可以用np.maximum.reduceat来得到输出，但这只适用于第一行。

indx = np.unique(data[2].astype(int),return_index=1)[1]
np.maximum.reduceat(data[0], indx)
   # not works for the whole `data`
# Output
array([0.66176217, 0.45114132])

我也尝试过创建一个掩码，但这并不总是能像例子那样工作。

mask = np.maximum.reduceat(data[0], indx)
data[:, np.isin(data[0], mask)]

# Output
array([[ 0.45114132,  0.66176217,  0.45114132],
       [11.        , 10.        ,  4.        ],
       [ 0.        ,  0.        ,  1.        ]])

另外，我想避免使用任何循环或列表的创建。
因为处理大量数据会太慢。

Answer 1

In [75]: data = np.array([[ 0.45114132,  0.31522008,  0.66176217,  0.45114132,  0.26872137],
    ...:        [11.        ,  6.        , 10.        ,  4.        ,  8.        ],
    ...:        [ 0.        ,  0.        ,  0.        ,  1.        ,  1.        ]])

In [76]: indx = np.unique(data[2].astype(int),return_index=1)[1]

In [77]: indx
Out[77]: array([0, 3], dtype=int64)

只需要提供整个数组和轴，就可以获取所有行的最大值：

In [78]: np.maximum.reduceat(data, indx, axis=1)
Out[78]: 
array([[ 0.66176217,  0.45114132],
       [11.        ,  8.        ],
       [ 0.        ,  1.        ]])

这会独立地对每一行应用最大值。

编辑

如果只想获取第一行的最大值，并返回整列，可以使用 split 将数组分成几个块，然后取每个块的最大值。

In [28]: datas = np.split(data, indx[1:], axis=1); datas
Out[28]: 
[array([[ 0.45114132,  0.31522008,  0.66176217],
        [11.        ,  6.        , 10.        ],
        [ 0.        ,  0.        ,  0.        ]]),
 array([[0.45114132, 0.26872137],
        [4.        , 8.        ],
        [1.        , 1.        ]])]

In [29]: [d[:,np.argmax(d[0])] for d in datas]
Out[29]: 
[array([ 0.66176217, 10.        ,  0.        ]),
 array([0.45114132, 4.        , 1.        ])]

In [30]: np.transpose(_)
Out[30]: 
array([[ 0.66176217,  0.45114132],
       [10.        ,  4.        ],
       [ 0.        ,  1.        ]])

是的，这确实在分块时有一个循环，但我不确定这比 reduceat 慢多少。它可能取决于块的数量与总列数的相对关系。一般来说，在复杂任务上进行几次迭代的性能还是不错的。

附注：argmax 没有 reduceat，所以不能像 maximum 那样使用。

Answer 2

我希望我理解了你的问题。你可以使用 np.lexsort 来根据 row[2] 的顺序对 row[0] 进行排序，然后用 np.unique 得到正确的值：

idx = np.unique(data[2].astype(int), return_index=1)[1]
mask = np.lexsort([-data[0], data[2]])[idx]
print(data[:, mask])

输出结果是：

[[ 0.66176217  0.45114132]
 [10.          4.        ]
 [ 0.          1.        ]]

Answer 3

使用 np.lexsort，你可以把所有的操作合并成一行代码。

data[:,np.lexsort(data[[0, 2]])[np.r_[data[2,:-1] != data[2,1:], True]]]

array([[ 0.66176217,  0.45114132],
       [10.        ,  4.        ],
       [ 0.        ,  1.        ]])