python之numpy.amax 中的一个键

在 Python 的标准中 max函数我可以传入一个key参数:

s = numpy.array(['one','two','three']) 
max(s) # 'two' (lexicographically last) 
max(s, key=len) # 'three' (longest string) 

对于更大的(多维)数组，我们不能再使用max，但我们可以使用numpy.amax。 ...不幸的是不提供key参数。

t = numpy.array([['one','two','three'], 
                 ['four','five','six']],  
                dtype='object') 
numpy.amax(t) # 'two` (max of the flat array) 
numpy.amax(t, axis=1) # array([two, six], dtype=object) (max of first row, followed by max of second row) 

我希望能够做的是:

amax2(t, key=len) # 'three' 
amax2(t, key=len, axis=1) # array([three, four], dtype=object) 

是否有内置方法来执行此操作？

注意:第一次尝试写这个问题时，我无法得到 amax working in this toy example !

请您参考如下方法:

这是一种非内置的方式(它在使用key时缺少features of amax的out和keepdim参数)，它看起来很长:

def amax2(x, *args, **kwargs): 
    if 'key' not in kwargs: 
        return numpy.amax(x,*args,**kwargs) 
    else: 
        key = kwargs.pop('key') # e.g. len, pop so no TypeError: unexpected keyword 
        x_key = numpy.vectorize(key)(x) # apply key to x element-wise 
        axis = kwargs.get('axis') # either None or axis is set in kwargs 
        if len(args)>=2: # axis is set in args 
            axis = args[1] 
 
        # The following is kept verbose, but could be made more efficient/shorter     
        if axis is None: # max of flattened 
            max_flat_index = numpy.argmax(x_key, axis=axis) 
            max_tuple_index = numpy.unravel_index(max_flat_index, x.shape) 
            return x[max_tuple_index] 
        elif axis == 0: # max in each column 
            max_indices = numpy.argmax(x_key, axis=axis) 
            return numpy.array( 
                 [ x[max_i, i] # reorder for col 
                     for i, max_i in enumerate(max_indices) ],  
                 dtype=x.dtype) 
        elif axis == 1: # max in each row 
            max_indices = numpy.argmax(x_key, axis=axis) 
            return numpy.array( 
                 [ x[i, max_i] 
                     for i, max_i in enumerate(max_indices) ], 
                 dtype=x.dtype) 

此功能的想法是从 @PeterSobot's answer 的第二部分扩展而来的到我之前的问题。